From 1414c9d0da946d33d72bd7ec7c2eb8efe4f03f38 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:00:23 +0200
Subject: [PATCH 1/7] Add torch stable ABI foundation: torch_compat layer and
 stable Tensor caster

Start of the incremental migration to the torch stable ABI on pybind11:
- csrc/torch_compat.h: dual-mode compat layer (stable shims with
  NVTE_WITH_TORCH_STABLE, full-ABI polyfills otherwise)
- pybind11 type_caster<torch::stable::Tensor> built on it
- makeTransformerEngineTensor/convertTorchShape overloads for stable tensors
- extensions/recipe.cpp ported to torch::stable::Tensor
- experimental NVTE_TORCH_STABLE_ABI build flag

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 build_tools/pytorch.py                        |  7 ++
 transformer_engine/pytorch/csrc/common.cpp    | 21 +++++
 transformer_engine/pytorch/csrc/common.h      |  6 ++
 transformer_engine/pytorch/csrc/extensions.h  | 11 ++-
 .../pytorch/csrc/extensions/recipe.cpp        | 28 +++----
 .../csrc/extensions/stable_tensor_caster.h    | 50 +++++++++++
 .../pytorch/csrc/torch_compat.h               | 82 +++++++++++++++++++
 7 files changed, 184 insertions(+), 21 deletions(-)
 create mode 100644 transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
 create mode 100644 transformer_engine/pytorch/csrc/torch_compat.h

diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
index 2bb238c522..bbd1f2c803 100644
--- a/build_tools/pytorch.py
+++ b/build_tools/pytorch.py
@@ -111,6 +111,13 @@ def setup_pytorch_extension(
     if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))):
         cxx_flags.append("-DNVTE_WITH_CUBLASMP")
 
+    # Experimental: build the torch_compat layer against the torch stable ABI
+    # (requires torch >= 2.14). Without the flag the same code compiles against
+    # the full torch ABI. See csrc/torch_compat.h.
+    if bool(int(os.getenv("NVTE_TORCH_STABLE_ABI", "0"))):
+        cxx_flags.append("-DNVTE_WITH_TORCH_STABLE")
+        cxx_flags.append("-DTORCH_TARGET_VERSION=0x020e000000000000")
+
     # Construct PyTorch CUDA extension
     sources = [str(path) for path in sources]
     include_dirs = [str(path) for path in include_dirs]
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index d85dcda159..782cb8eb1c 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -62,6 +62,19 @@ NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape) {
   return ret;
 }
 
+NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape) {
+  NVTEShape ret;
+  ret.ndim = torch_shape.size();
+  constexpr int max_dimensions = sizeof(ret.data) / sizeof(size_t);
+  NVTE_CHECK(ret.ndim < max_dimensions,
+             "Torch tensor has too many dimensions. Max supported: ", max_dimensions, " and got ",
+             ret.ndim, ".");
+  for (size_t i = 0; i < ret.ndim; ++i) {
+    ret.data[i] = static_cast<size_t>(torch_shape[i]);
+  }
+  return ret;
+}
+
 std::unique_ptr<Quantizer> convert_quantizer(py::handle quantizer) {
   init_extension();
   if (quantizer.is_none()) {
@@ -165,6 +178,14 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor)
   return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype);
 }
 
+transformer_engine::TensorWrapper makeTransformerEngineTensor(
+    const torch::stable::Tensor& tensor) {
+  transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type());
+  const auto sizes = tensor.sizes();
+  std::vector<size_t> shape(sizes.begin(), sizes.end());
+  return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype);
+}
+
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
            std::vector<NVTETensor*>, size_t, size_t>
 makeTransformerEngineTensorList(std::vector<std::vector<at::Tensor>> at_tensor_lists) {
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
index aa0e0c87fe..3c43f8c6da 100644
--- a/transformer_engine/pytorch/csrc/common.h
+++ b/transformer_engine/pytorch/csrc/common.h
@@ -56,6 +56,8 @@
 #include "c10/util/ArrayRef.h"
 #include "common/util/logging.h"
 #include "extensions/pybind_dtype_caster.h"
+#include "extensions/stable_tensor_caster.h"
+#include "torch_compat.h"
 
 namespace transformer_engine::pytorch {
 
@@ -544,6 +546,8 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(void* data_ptr,
 
 transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor);
 
+transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch::stable::Tensor& tensor);
+
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
            std::vector<NVTETensor*>, size_t, size_t>
 makeTransformerEngineTensorList(std::vector<std::vector<at::Tensor>> at_tensor_lists);
@@ -581,6 +585,8 @@ size_t ceildiv(size_t numer, size_t denom);
 
 NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape);
 
+NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape);
+
 std::vector<size_t> convert_shape_back_from_fp4(const std::vector<size_t>& shape, bool transpose);
 
 // Flatten an N-D shape to 2D: {product(shape[:-1]), shape[-1]}.
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 95f02c64f0..b498613cdb 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -437,13 +437,12 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads
  * FP8 recipe
  **************************************************************************************************/
 
-void compute_amax(const at::Tensor &tensor, at::Tensor &amax);
+void compute_amax(const torch::stable::Tensor &tensor, torch::stable::Tensor &amax);
 
-void fused_amax_and_scale_update_after_reduction(const at::Tensor &amax_reduction_buffer,
-                                                 std::vector<at::Tensor> amax_histories,
-                                                 std::vector<at::Tensor> scales,
-                                                 const std::string &amax_compute_algo,
-                                                 DType fp8_dtype, float margin);
+void fused_amax_and_scale_update_after_reduction(
+    const torch::stable::Tensor &amax_reduction_buffer,
+    std::vector<torch::stable::Tensor> amax_histories, std::vector<torch::stable::Tensor> scales,
+    const std::string &amax_compute_algo, DType fp8_dtype, float margin);
 
 // Note that the start_offset is the logical offset along the tensor dimension.
 // The offset in bytes is start_offset * sizeof(tensor.dtype)
diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
index 9be288d2f7..08e25072bd 100644
--- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
@@ -4,36 +4,34 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
 #include <string>
 
 #include "../extensions.h"
+#include "../torch_compat.h"
 #include "transformer_engine/transformer_engine.h"
 
 namespace transformer_engine::pytorch {
 
-void compute_amax(const at::Tensor& tensor, at::Tensor& amax) {
-  auto input_tensor = tensor.contiguous();
+void compute_amax(const torch::stable::Tensor& tensor, torch::stable::Tensor& amax) {
+  auto input_tensor = torch::stable::contiguous(tensor);
   const TensorWrapper& te_input = makeTransformerEngineTensor(input_tensor);
 
-  TORCH_CHECK(amax.scalar_type() == at::kFloat, "amax must be a float tensor");
-  TORCH_CHECK(amax.numel() == 1, "amax must have exactly one element");
-  auto* amax_ptr = amax.data_ptr<float>();
+  NVTE_CHECK(amax.scalar_type() == torch::headeronly::ScalarType::Float,
+             "amax must be a float tensor");
+  NVTE_CHECK(amax.numel() == 1, "amax must have exactly one element");
+  auto* amax_ptr = static_cast<float*>(amax.data_ptr());
   TensorWrapper fake_te_output(
       /*dptr=*/nullptr, te_input.shape(),
       DType::kFloat32,  // It doesn't matter because we only compute amax.
       amax_ptr);
 
-  nvte_compute_amax(te_input.data(), fake_te_output.data(), at::cuda::getCurrentCUDAStream());
+  nvte_compute_amax(te_input.data(), fake_te_output.data(), torch_compat::getCurrentCUDAStream());
 }
 
-void fused_amax_and_scale_update_after_reduction(const at::Tensor& amax_reduction_buffer,
-                                                 std::vector<at::Tensor> amax_histories,
-                                                 std::vector<at::Tensor> scales,
-                                                 const std::string& amax_compute_algo,
-                                                 DType fp8_dtype, float margin) {
+void fused_amax_and_scale_update_after_reduction(
+    const torch::stable::Tensor& amax_reduction_buffer,
+    std::vector<torch::stable::Tensor> amax_histories, std::vector<torch::stable::Tensor> scales,
+    const std::string& amax_compute_algo, DType fp8_dtype, float margin) {
   size_t num_tensors = amax_histories.size();
 
   // Allocate amax history and scale NVTETensors as batches
@@ -58,7 +56,7 @@ void fused_amax_and_scale_update_after_reduction(const at::Tensor& amax_reductio
       makeTransformerEngineTensor(amax_reduction_buffer).data(),
       std::vector<NVTETensor>(te_amax_histories.begin(), te_amax_histories.end()),
       std::vector<NVTETensor>(te_scales.begin(), te_scales.end()), amax_compute_algo.c_str(),
-      static_cast<NVTEDType>(fp8_dtype), margin, at::cuda::getCurrentCUDAStream());
+      static_cast<NVTEDType>(fp8_dtype), margin, torch_compat::getCurrentCUDAStream());
 }
 
 }  // namespace transformer_engine::pytorch
diff --git a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
new file mode 100644
index 0000000000..e289207abc
--- /dev/null
+++ b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
@@ -0,0 +1,50 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
+#define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
+
+#include <pybind11/pybind11.h>
+
+#include "../torch_compat.h"
+
+namespace pybind11 {
+namespace detail {
+
+/*! @brief Custom type caster for ``torch::stable::Tensor``.
+ *
+ * Lets pybind-bound functions take/return ``torch::stable::Tensor`` directly:
+ * a ``torch.Tensor`` argument is unwrapped into a stable tensor sharing the
+ * same TensorImpl, and a returned stable tensor is wrapped back into a
+ * ``torch.Tensor``.
+ *
+ * NOTE: As a compile-time specialization this must be visible in every
+ * translation unit that converts ``torch::stable::Tensor`` (it is pulled in
+ * via the PyTorch extension's ``common.h``), otherwise different TUs would
+ * instantiate different casters for the same type (ODR violation).
+ */
+template <>
+struct type_caster<torch::stable::Tensor> {
+ public:
+  PYBIND11_TYPE_CASTER(torch::stable::Tensor, const_name("torch.Tensor"));
+
+  bool load(handle src, bool) {
+    if (!src || !transformer_engine::pytorch::torch_compat::is_tensor_pyobject(src.ptr())) {
+      return false;
+    }
+    value = transformer_engine::pytorch::torch_compat::tensor_from_pyobject(src.ptr());
+    return true;
+  }
+
+  static handle cast(const torch::stable::Tensor &src, return_value_policy, handle) {
+    return handle(transformer_engine::pytorch::torch_compat::tensor_to_pyobject(src));
+  }
+};
+
+}  // namespace detail
+}  // namespace pybind11
+
+#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
diff --git a/transformer_engine/pytorch/csrc/torch_compat.h b/transformer_engine/pytorch/csrc/torch_compat.h
new file mode 100644
index 0000000000..6f0d365b9c
--- /dev/null
+++ b/transformer_engine/pytorch/csrc/torch_compat.h
@@ -0,0 +1,82 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
+#define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
+
+#include <Python.h>
+#include <cuda_runtime.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+
+#ifdef NVTE_WITH_TORCH_STABLE
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/pyobject.h>
+#else
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/csrc/autograd/python_variable.h>
+#include <torch/csrc/inductor/aoti_torch/utils.h>
+#endif
+
+#include "common/util/logging.h"
+
+/* Compatibility layer for the incremental migration to the torch stable ABI.
+ * Code is written against the torch::stable surface. With NVTE_WITH_TORCH_STABLE
+ * these helpers forward to the stable shims (torch >= 2.14); without it the same
+ * functionality is polyfilled with the full torch ABI, so the non-stable build
+ * keeps supporting older torch versions. */
+namespace transformer_engine::pytorch::torch_compat {
+
+inline bool is_tensor_pyobject(PyObject *obj) {
+#ifdef NVTE_WITH_TORCH_STABLE
+  static PyObject *tensor_type = [] {
+    PyObject *mod = PyImport_ImportModule("torch");
+    NVTE_CHECK(mod != nullptr, "Could not import torch");
+    PyObject *type = PyObject_GetAttrString(mod, "Tensor");
+    Py_DECREF(mod);
+    NVTE_CHECK(type != nullptr, "Could not get torch.Tensor");
+    return type;
+  }();
+  return PyObject_IsInstance(obj, tensor_type) == 1;
+#else
+  return THPVariable_Check(obj);
+#endif
+}
+
+/* Borrowed torch.Tensor PyObject -> stable Tensor sharing the TensorImpl.
+ * The GIL must be held. */
+inline torch::stable::Tensor tensor_from_pyobject(PyObject *obj) {
+#ifdef NVTE_WITH_TORCH_STABLE
+  return torch::stable::tensor_from_pyobject(obj);
+#else
+  return torch::stable::Tensor(
+      torch::aot_inductor::new_tensor_handle(at::Tensor(THPVariable_Unpack(obj))));
+#endif
+}
+
+/* Stable Tensor -> new-reference torch.Tensor PyObject. The GIL must be held. */
+inline PyObject *tensor_to_pyobject(const torch::stable::Tensor &tensor) {
+#ifdef NVTE_WITH_TORCH_STABLE
+  return static_cast<PyObject *>(torch::stable::tensor_to_pyobject(tensor));
+#else
+  return THPVariable_Wrap(*torch::aot_inductor::tensor_handle_to_tensor_pointer(tensor.get()));
+#endif
+}
+
+inline cudaStream_t getCurrentCUDAStream() {
+#ifdef NVTE_WITH_TORCH_STABLE
+  return static_cast<cudaStream_t>(
+      torch::stable::accelerator::getCurrentStream(
+          torch::stable::accelerator::getCurrentDeviceIndex())
+          .nativeHandle());
+#else
+  return at::cuda::getCurrentCUDAStream();
+#endif
+}
+
+}  // namespace transformer_engine::pytorch::torch_compat
+
+#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_

From ec32413d323c1b4ff3058aaed06e8d43d891ee77 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:10:42 +0200
Subject: [PATCH 2/7] Keep torch>=2.1 support: default build no longer touches
 stable headers

torch_compat::Tensor is at::Tensor in the default build (torch's own pybind
caster applies) and torch::stable::Tensor only under NVTE_WITH_TORCH_STABLE;
the stable tensor caster and the stable overloads are compiled only in stable
mode. Migrated code targets the torch_compat surface instead of torch::stable
directly.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/csrc/common.cpp    |  7 ++-
 transformer_engine/pytorch/csrc/common.h      |  6 ++-
 transformer_engine/pytorch/csrc/extensions.h  |  6 +--
 .../pytorch/csrc/extensions/recipe.cpp        | 10 ++--
 .../csrc/extensions/stable_tensor_caster.h    | 26 ++++++++--
 .../pytorch/csrc/torch_compat.h               | 52 +++++--------------
 6 files changed, 55 insertions(+), 52 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index 782cb8eb1c..e8f191cbe8 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -62,6 +62,7 @@ NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape) {
   return ret;
 }
 
+#ifdef NVTE_WITH_TORCH_STABLE
 NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape) {
   NVTEShape ret;
   ret.ndim = torch_shape.size();
@@ -74,6 +75,7 @@ NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch
   }
   return ret;
 }
+#endif
 
 std::unique_ptr<Quantizer> convert_quantizer(py::handle quantizer) {
   init_extension();
@@ -178,13 +180,14 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor)
   return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype);
 }
 
-transformer_engine::TensorWrapper makeTransformerEngineTensor(
-    const torch::stable::Tensor& tensor) {
+#ifdef NVTE_WITH_TORCH_STABLE
+transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_compat::Tensor& tensor) {
   transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type());
   const auto sizes = tensor.sizes();
   std::vector<size_t> shape(sizes.begin(), sizes.end());
   return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype);
 }
+#endif
 
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
            std::vector<NVTETensor*>, size_t, size_t>
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
index 3c43f8c6da..500546784b 100644
--- a/transformer_engine/pytorch/csrc/common.h
+++ b/transformer_engine/pytorch/csrc/common.h
@@ -546,7 +546,9 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(void* data_ptr,
 
 transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor);
 
-transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch::stable::Tensor& tensor);
+#ifdef NVTE_WITH_TORCH_STABLE
+transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_compat::Tensor& tensor);
+#endif
 
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
            std::vector<NVTETensor*>, size_t, size_t>
@@ -585,7 +587,9 @@ size_t ceildiv(size_t numer, size_t denom);
 
 NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape);
 
+#ifdef NVTE_WITH_TORCH_STABLE
 NVTEShape convertTorchShape(const torch::headeronly::IntHeaderOnlyArrayRef torch_shape);
+#endif
 
 std::vector<size_t> convert_shape_back_from_fp4(const std::vector<size_t>& shape, bool transpose);
 
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index b498613cdb..1caec52a40 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -437,11 +437,11 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads
  * FP8 recipe
  **************************************************************************************************/
 
-void compute_amax(const torch::stable::Tensor &tensor, torch::stable::Tensor &amax);
+void compute_amax(const torch_compat::Tensor &tensor, torch_compat::Tensor &amax);
 
 void fused_amax_and_scale_update_after_reduction(
-    const torch::stable::Tensor &amax_reduction_buffer,
-    std::vector<torch::stable::Tensor> amax_histories, std::vector<torch::stable::Tensor> scales,
+    const torch_compat::Tensor &amax_reduction_buffer,
+    std::vector<torch_compat::Tensor> amax_histories, std::vector<torch_compat::Tensor> scales,
     const std::string &amax_compute_algo, DType fp8_dtype, float margin);
 
 // Note that the start_offset is the logical offset along the tensor dimension.
diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
index 08e25072bd..513244c087 100644
--- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
@@ -12,11 +12,11 @@
 
 namespace transformer_engine::pytorch {
 
-void compute_amax(const torch::stable::Tensor& tensor, torch::stable::Tensor& amax) {
-  auto input_tensor = torch::stable::contiguous(tensor);
+void compute_amax(const torch_compat::Tensor& tensor, torch_compat::Tensor& amax) {
+  auto input_tensor = torch_compat::contiguous(tensor);
   const TensorWrapper& te_input = makeTransformerEngineTensor(input_tensor);
 
-  NVTE_CHECK(amax.scalar_type() == torch::headeronly::ScalarType::Float,
+  NVTE_CHECK(amax.scalar_type() == torch_compat::ScalarType::Float,
              "amax must be a float tensor");
   NVTE_CHECK(amax.numel() == 1, "amax must have exactly one element");
   auto* amax_ptr = static_cast<float*>(amax.data_ptr());
@@ -29,8 +29,8 @@ void compute_amax(const torch::stable::Tensor& tensor, torch::stable::Tensor& am
 }
 
 void fused_amax_and_scale_update_after_reduction(
-    const torch::stable::Tensor& amax_reduction_buffer,
-    std::vector<torch::stable::Tensor> amax_histories, std::vector<torch::stable::Tensor> scales,
+    const torch_compat::Tensor& amax_reduction_buffer,
+    std::vector<torch_compat::Tensor> amax_histories, std::vector<torch_compat::Tensor> scales,
     const std::string& amax_compute_algo, DType fp8_dtype, float margin) {
   size_t num_tensors = amax_histories.size();
 
diff --git a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
index e289207abc..3c42f51132 100644
--- a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
+++ b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
@@ -7,9 +7,15 @@
 #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
 #define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
 
+/* In the default (non-stable) build torch_compat::Tensor is at::Tensor and
+ * torch's own pybind caster applies; this caster exists only in stable mode. */
+#ifdef NVTE_WITH_TORCH_STABLE
+
+#include <Python.h>
 #include <pybind11/pybind11.h>
 
 #include "../torch_compat.h"
+#include "common/util/logging.h"
 
 namespace pybind11 {
 namespace detail {
@@ -31,20 +37,34 @@ struct type_caster<torch::stable::Tensor> {
  public:
   PYBIND11_TYPE_CASTER(torch::stable::Tensor, const_name("torch.Tensor"));
 
+  static bool is_tensor(PyObject *obj) {
+    static PyObject *tensor_type = [] {
+      PyObject *mod = PyImport_ImportModule("torch");
+      NVTE_CHECK(mod != nullptr, "Could not import torch");
+      PyObject *type = PyObject_GetAttrString(mod, "Tensor");
+      Py_DECREF(mod);
+      NVTE_CHECK(type != nullptr, "Could not get torch.Tensor");
+      return type;
+    }();
+    return PyObject_IsInstance(obj, tensor_type) == 1;
+  }
+
   bool load(handle src, bool) {
-    if (!src || !transformer_engine::pytorch::torch_compat::is_tensor_pyobject(src.ptr())) {
+    if (!src || !is_tensor(src.ptr())) {
       return false;
     }
-    value = transformer_engine::pytorch::torch_compat::tensor_from_pyobject(src.ptr());
+    value = torch::stable::tensor_from_pyobject(src.ptr());
     return true;
   }
 
   static handle cast(const torch::stable::Tensor &src, return_value_policy, handle) {
-    return handle(transformer_engine::pytorch::torch_compat::tensor_to_pyobject(src));
+    return handle(static_cast<PyObject *>(torch::stable::tensor_to_pyobject(src)));
   }
 };
 
 }  // namespace detail
 }  // namespace pybind11
 
+#endif  // NVTE_WITH_TORCH_STABLE
+
 #endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
diff --git a/transformer_engine/pytorch/csrc/torch_compat.h b/transformer_engine/pytorch/csrc/torch_compat.h
index 6f0d365b9c..5046b12a19 100644
--- a/transformer_engine/pytorch/csrc/torch_compat.h
+++ b/transformer_engine/pytorch/csrc/torch_compat.h
@@ -7,62 +7,38 @@
 #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
 #define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
 
-#include <Python.h>
 #include <cuda_runtime.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/csrc/stable/tensor.h>
 
 #ifdef NVTE_WITH_TORCH_STABLE
 #include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
 #include <torch/csrc/stable/pyobject.h>
+#include <torch/csrc/stable/tensor.h>
 #else
+#include <ATen/ATen.h>
 #include <ATen/cuda/CUDAContext.h>
-#include <torch/csrc/autograd/python_variable.h>
-#include <torch/csrc/inductor/aoti_torch/utils.h>
 #endif
 
-#include "common/util/logging.h"
-
 /* Compatibility layer for the incremental migration to the torch stable ABI.
- * Code is written against the torch::stable surface. With NVTE_WITH_TORCH_STABLE
- * these helpers forward to the stable shims (torch >= 2.14); without it the same
- * functionality is polyfilled with the full torch ABI, so the non-stable build
- * keeps supporting older torch versions. */
+ * Migrated code is written against this surface, which is restricted to what
+ * torch::stable provides. With NVTE_WITH_TORCH_STABLE it maps to torch::stable
+ * (requires torch >= 2.14); without it (the default) it maps to the full torch
+ * ABI, keeping support for older torch versions intact. */
 namespace transformer_engine::pytorch::torch_compat {
 
-inline bool is_tensor_pyobject(PyObject *obj) {
 #ifdef NVTE_WITH_TORCH_STABLE
-  static PyObject *tensor_type = [] {
-    PyObject *mod = PyImport_ImportModule("torch");
-    NVTE_CHECK(mod != nullptr, "Could not import torch");
-    PyObject *type = PyObject_GetAttrString(mod, "Tensor");
-    Py_DECREF(mod);
-    NVTE_CHECK(type != nullptr, "Could not get torch.Tensor");
-    return type;
-  }();
-  return PyObject_IsInstance(obj, tensor_type) == 1;
+using Tensor = torch::stable::Tensor;
+using ScalarType = torch::headeronly::ScalarType;
 #else
-  return THPVariable_Check(obj);
+using Tensor = at::Tensor;
+using ScalarType = at::ScalarType;
 #endif
-}
-
-/* Borrowed torch.Tensor PyObject -> stable Tensor sharing the TensorImpl.
- * The GIL must be held. */
-inline torch::stable::Tensor tensor_from_pyobject(PyObject *obj) {
-#ifdef NVTE_WITH_TORCH_STABLE
-  return torch::stable::tensor_from_pyobject(obj);
-#else
-  return torch::stable::Tensor(
-      torch::aot_inductor::new_tensor_handle(at::Tensor(THPVariable_Unpack(obj))));
-#endif
-}
 
-/* Stable Tensor -> new-reference torch.Tensor PyObject. The GIL must be held. */
-inline PyObject *tensor_to_pyobject(const torch::stable::Tensor &tensor) {
+inline Tensor contiguous(const Tensor &tensor) {
 #ifdef NVTE_WITH_TORCH_STABLE
-  return static_cast<PyObject *>(torch::stable::tensor_to_pyobject(tensor));
+  return torch::stable::contiguous(tensor);
 #else
-  return THPVariable_Wrap(*torch::aot_inductor::tensor_handle_to_tensor_pointer(tensor.get()));
+  return tensor.contiguous();
 #endif
 }
 

From 783f0b35c5b032f3d6de2388c4912687e60ec799 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:51:24 +0200
Subject: [PATCH 3/7] Rename torch_compat to torch_stable

The namespace is the torch::stable API surface TE migrates to, not a
backwards-compat shim; the name states the contract: ported files use
torch_stable::, never at::.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/csrc/common.cpp       |  2 +-
 transformer_engine/pytorch/csrc/common.h         |  4 ++--
 transformer_engine/pytorch/csrc/extensions.h     |  6 +++---
 .../pytorch/csrc/extensions/recipe.cpp           | 16 ++++++++--------
 .../csrc/{torch_compat.h => torch_stable.h}      | 10 +++++-----
 5 files changed, 19 insertions(+), 19 deletions(-)
 rename transformer_engine/pytorch/csrc/{torch_compat.h => torch_stable.h} (85%)

diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index e8f191cbe8..114e11e275 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -181,7 +181,7 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor)
 }
 
 #ifdef NVTE_WITH_TORCH_STABLE
-transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_compat::Tensor& tensor) {
+transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_stable::Tensor& tensor) {
   transformer_engine::DType dtype = GetTransformerEngineDType(tensor.scalar_type());
   const auto sizes = tensor.sizes();
   std::vector<size_t> shape(sizes.begin(), sizes.end());
diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h
index 500546784b..3d33d2ac11 100644
--- a/transformer_engine/pytorch/csrc/common.h
+++ b/transformer_engine/pytorch/csrc/common.h
@@ -57,7 +57,7 @@
 #include "common/util/logging.h"
 #include "extensions/pybind_dtype_caster.h"
 #include "extensions/stable_tensor_caster.h"
-#include "torch_compat.h"
+#include "torch_stable.h"
 
 namespace transformer_engine::pytorch {
 
@@ -547,7 +547,7 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(void* data_ptr,
 transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor);
 
 #ifdef NVTE_WITH_TORCH_STABLE
-transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_compat::Tensor& tensor);
+transformer_engine::TensorWrapper makeTransformerEngineTensor(const torch_stable::Tensor& tensor);
 #endif
 
 std::tuple<std::vector<transformer_engine::TensorWrapper>, std::vector<std::vector<NVTETensor>>,
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 1caec52a40..848c57f3f8 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -437,11 +437,11 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads
  * FP8 recipe
  **************************************************************************************************/
 
-void compute_amax(const torch_compat::Tensor &tensor, torch_compat::Tensor &amax);
+void compute_amax(const torch_stable::Tensor &tensor, torch_stable::Tensor &amax);
 
 void fused_amax_and_scale_update_after_reduction(
-    const torch_compat::Tensor &amax_reduction_buffer,
-    std::vector<torch_compat::Tensor> amax_histories, std::vector<torch_compat::Tensor> scales,
+    const torch_stable::Tensor &amax_reduction_buffer,
+    std::vector<torch_stable::Tensor> amax_histories, std::vector<torch_stable::Tensor> scales,
     const std::string &amax_compute_algo, DType fp8_dtype, float margin);
 
 // Note that the start_offset is the logical offset along the tensor dimension.
diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
index 513244c087..70e3fa1693 100644
--- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
@@ -7,16 +7,16 @@
 #include <string>
 
 #include "../extensions.h"
-#include "../torch_compat.h"
+#include "../torch_stable.h"
 #include "transformer_engine/transformer_engine.h"
 
 namespace transformer_engine::pytorch {
 
-void compute_amax(const torch_compat::Tensor& tensor, torch_compat::Tensor& amax) {
-  auto input_tensor = torch_compat::contiguous(tensor);
+void compute_amax(const torch_stable::Tensor& tensor, torch_stable::Tensor& amax) {
+  auto input_tensor = torch_stable::contiguous(tensor);
   const TensorWrapper& te_input = makeTransformerEngineTensor(input_tensor);
 
-  NVTE_CHECK(amax.scalar_type() == torch_compat::ScalarType::Float,
+  NVTE_CHECK(amax.scalar_type() == torch_stable::ScalarType::Float,
              "amax must be a float tensor");
   NVTE_CHECK(amax.numel() == 1, "amax must have exactly one element");
   auto* amax_ptr = static_cast<float*>(amax.data_ptr());
@@ -25,12 +25,12 @@ void compute_amax(const torch_compat::Tensor& tensor, torch_compat::Tensor& amax
       DType::kFloat32,  // It doesn't matter because we only compute amax.
       amax_ptr);
 
-  nvte_compute_amax(te_input.data(), fake_te_output.data(), torch_compat::getCurrentCUDAStream());
+  nvte_compute_amax(te_input.data(), fake_te_output.data(), torch_stable::getCurrentCUDAStream());
 }
 
 void fused_amax_and_scale_update_after_reduction(
-    const torch_compat::Tensor& amax_reduction_buffer,
-    std::vector<torch_compat::Tensor> amax_histories, std::vector<torch_compat::Tensor> scales,
+    const torch_stable::Tensor& amax_reduction_buffer,
+    std::vector<torch_stable::Tensor> amax_histories, std::vector<torch_stable::Tensor> scales,
     const std::string& amax_compute_algo, DType fp8_dtype, float margin) {
   size_t num_tensors = amax_histories.size();
 
@@ -56,7 +56,7 @@ void fused_amax_and_scale_update_after_reduction(
       makeTransformerEngineTensor(amax_reduction_buffer).data(),
       std::vector<NVTETensor>(te_amax_histories.begin(), te_amax_histories.end()),
       std::vector<NVTETensor>(te_scales.begin(), te_scales.end()), amax_compute_algo.c_str(),
-      static_cast<NVTEDType>(fp8_dtype), margin, torch_compat::getCurrentCUDAStream());
+      static_cast<NVTEDType>(fp8_dtype), margin, torch_stable::getCurrentCUDAStream());
 }
 
 }  // namespace transformer_engine::pytorch
diff --git a/transformer_engine/pytorch/csrc/torch_compat.h b/transformer_engine/pytorch/csrc/torch_stable.h
similarity index 85%
rename from transformer_engine/pytorch/csrc/torch_compat.h
rename to transformer_engine/pytorch/csrc/torch_stable.h
index 5046b12a19..19a15b7308 100644
--- a/transformer_engine/pytorch/csrc/torch_compat.h
+++ b/transformer_engine/pytorch/csrc/torch_stable.h
@@ -4,8 +4,8 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
-#define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
+#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_
+#define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_
 
 #include <cuda_runtime.h>
 
@@ -24,7 +24,7 @@
  * torch::stable provides. With NVTE_WITH_TORCH_STABLE it maps to torch::stable
  * (requires torch >= 2.14); without it (the default) it maps to the full torch
  * ABI, keeping support for older torch versions intact. */
-namespace transformer_engine::pytorch::torch_compat {
+namespace transformer_engine::pytorch::torch_stable {
 
 #ifdef NVTE_WITH_TORCH_STABLE
 using Tensor = torch::stable::Tensor;
@@ -53,6 +53,6 @@ inline cudaStream_t getCurrentCUDAStream() {
 #endif
 }
 
-}  // namespace transformer_engine::pytorch::torch_compat
+}  // namespace transformer_engine::pytorch::torch_stable
 
-#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_COMPAT_H_
+#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_

From 98c7864d81f64adb0485bd145087a1a6b95a258b Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:51:25 +0200
Subject: [PATCH 4/7] Use torch::stable::is_tensor_pyobject in the stable
 tensor caster

Replaces the import-torch + isinstance probe (~200ns/call, Python round
trip) with the C shim (~20ns). Requires the torch PR adding
torch_is_tensor_pyobject on top of torch 2.14.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../csrc/extensions/stable_tensor_caster.h    | 20 +++----------------
 1 file changed, 3 insertions(+), 17 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
index 3c42f51132..b4f6ff82be 100644
--- a/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
+++ b/transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h
@@ -7,15 +7,13 @@
 #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
 #define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_STABLE_TENSOR_CASTER_H_
 
-/* In the default (non-stable) build torch_compat::Tensor is at::Tensor and
+/* In the default (non-stable) build torch_stable::Tensor is at::Tensor and
  * torch's own pybind caster applies; this caster exists only in stable mode. */
 #ifdef NVTE_WITH_TORCH_STABLE
 
-#include <Python.h>
 #include <pybind11/pybind11.h>
 
-#include "../torch_compat.h"
-#include "common/util/logging.h"
+#include "../torch_stable.h"
 
 namespace pybind11 {
 namespace detail {
@@ -37,20 +35,8 @@ struct type_caster<torch::stable::Tensor> {
  public:
   PYBIND11_TYPE_CASTER(torch::stable::Tensor, const_name("torch.Tensor"));
 
-  static bool is_tensor(PyObject *obj) {
-    static PyObject *tensor_type = [] {
-      PyObject *mod = PyImport_ImportModule("torch");
-      NVTE_CHECK(mod != nullptr, "Could not import torch");
-      PyObject *type = PyObject_GetAttrString(mod, "Tensor");
-      Py_DECREF(mod);
-      NVTE_CHECK(type != nullptr, "Could not get torch.Tensor");
-      return type;
-    }();
-    return PyObject_IsInstance(obj, tensor_type) == 1;
-  }
-
   bool load(handle src, bool) {
-    if (!src || !is_tensor(src.ptr())) {
+    if (!src || !torch::stable::is_tensor_pyobject(src.ptr())) {
       return false;
     }
     value = torch::stable::tensor_from_pyobject(src.ptr());

From cb3933b12795e5bfa96a7c73fb144d7fd30392b7 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:51:25 +0200
Subject: [PATCH 5/7] Reject too-old runtime torch in stable-ABI builds

Stable-ABI builds link shims that only exist in libtorch >= 2.14; on an
older runtime the import would die with a raw dynamic-linker error. The
build now records TORCH_STABLE_ABI in a generated _build_config.py and
__init__.py fails with a clear message before loading the extension.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .gitignore                             |  1 +
 build_tools/pytorch.py                 | 17 ++++++++++++++---
 transformer_engine/pytorch/__init__.py | 17 +++++++++++++++++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 8a627a7e76..a467b895a0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -43,3 +43,4 @@ tensor_dumps/
 artifacts/
 .DS_Store
 .claude/
+transformer_engine/pytorch/_build_config.py
diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
index bbd1f2c803..e6fac85533 100644
--- a/build_tools/pytorch.py
+++ b/build_tools/pytorch.py
@@ -111,13 +111,24 @@ def setup_pytorch_extension(
     if bool(int(os.getenv("NVTE_WITH_CUBLASMP", 0))):
         cxx_flags.append("-DNVTE_WITH_CUBLASMP")
 
-    # Experimental: build the torch_compat layer against the torch stable ABI
+    # Experimental: build the torch_stable layer against the torch stable ABI
     # (requires torch >= 2.14). Without the flag the same code compiles against
-    # the full torch ABI. See csrc/torch_compat.h.
-    if bool(int(os.getenv("NVTE_TORCH_STABLE_ABI", "0"))):
+    # the full torch ABI. See csrc/torch_stable.h.
+    torch_stable_abi = bool(int(os.getenv("NVTE_TORCH_STABLE_ABI", "0")))
+    if torch_stable_abi:
         cxx_flags.append("-DNVTE_WITH_TORCH_STABLE")
         cxx_flags.append("-DTORCH_TARGET_VERSION=0x020e000000000000")
 
+    # Record build configuration for runtime checks (stable-ABI builds link
+    # shims that only exist in libtorch >= TORCH_TARGET_VERSION, so the Python
+    # package must reject older runtime torch before loading the extension).
+    build_config = Path(csrc_header_files).parent / "_build_config.py"
+    build_config.write_text(
+        "# Generated by build_tools/pytorch.py. Do not edit.\n"
+        f"TORCH_STABLE_ABI = {torch_stable_abi}\n"
+        "TORCH_STABLE_ABI_MIN_TORCH = (2, 14)\n"
+    )
+
     # Construct PyTorch CUDA extension
     sources = [str(path) for path in sources]
     include_dirs = [str(path) for path in include_dirs]
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 4c6b7fc67a..15b1069ca7 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -15,6 +15,23 @@
 
 assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}."
 
+try:
+    from transformer_engine.pytorch._build_config import (
+        TORCH_STABLE_ABI,
+        TORCH_STABLE_ABI_MIN_TORCH,
+    )
+except ImportError:
+    TORCH_STABLE_ABI = False
+    TORCH_STABLE_ABI_MIN_TORCH = None
+
+if TORCH_STABLE_ABI:
+    # Stable-ABI builds link shims that only exist in newer libtorch; loading
+    # them on an older runtime would fail with a raw dynamic-linker error.
+    assert torch_version() >= TORCH_STABLE_ABI_MIN_TORCH, (
+        "This Transformer Engine build uses the torch stable ABI and requires torch >="
+        f" {'.'.join(map(str, TORCH_STABLE_ABI_MIN_TORCH))} at runtime. Found {torch_version()}."
+    )
+
 load_framework_extension("torch")
 from transformer_engine.pytorch import constants
 from transformer_engine.pytorch.constants import DType

From 9e1809e519ae50f1375ddc8ae21e6dc1b7c20a3a Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 10 Aug 2026 11:55:20 +0200
Subject: [PATCH 6/7] Add module docstring to generated _build_config.py
 (pylint C0114)

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 build_tools/pytorch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py
index e6fac85533..0fd111d07b 100644
--- a/build_tools/pytorch.py
+++ b/build_tools/pytorch.py
@@ -124,7 +124,7 @@ def setup_pytorch_extension(
     # package must reject older runtime torch before loading the extension).
     build_config = Path(csrc_header_files).parent / "_build_config.py"
     build_config.write_text(
-        "# Generated by build_tools/pytorch.py. Do not edit.\n"
+        '"""Build configuration. Generated by build_tools/pytorch.py, do not edit."""\n'
         f"TORCH_STABLE_ABI = {torch_stable_abi}\n"
         "TORCH_STABLE_ABI_MIN_TORCH = (2, 14)\n"
     )

From f9d6424ac708f1277b9ea51e47131147849dfb8f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 10 Aug 2026 10:06:03 +0000
Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/csrc/extensions.h         |  9 +++++----
 .../pytorch/csrc/extensions/recipe.cpp               | 12 ++++++------
 transformer_engine/pytorch/csrc/torch_stable.h       |  7 +++----
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 848c57f3f8..b6c58593b3 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -439,10 +439,11 @@ at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads
 
 void compute_amax(const torch_stable::Tensor &tensor, torch_stable::Tensor &amax);
 
-void fused_amax_and_scale_update_after_reduction(
-    const torch_stable::Tensor &amax_reduction_buffer,
-    std::vector<torch_stable::Tensor> amax_histories, std::vector<torch_stable::Tensor> scales,
-    const std::string &amax_compute_algo, DType fp8_dtype, float margin);
+void fused_amax_and_scale_update_after_reduction(const torch_stable::Tensor &amax_reduction_buffer,
+                                                 std::vector<torch_stable::Tensor> amax_histories,
+                                                 std::vector<torch_stable::Tensor> scales,
+                                                 const std::string &amax_compute_algo,
+                                                 DType fp8_dtype, float margin);
 
 // Note that the start_offset is the logical offset along the tensor dimension.
 // The offset in bytes is start_offset * sizeof(tensor.dtype)
diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
index 70e3fa1693..158dcbfa70 100644
--- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp
@@ -16,8 +16,7 @@ void compute_amax(const torch_stable::Tensor& tensor, torch_stable::Tensor& amax
   auto input_tensor = torch_stable::contiguous(tensor);
   const TensorWrapper& te_input = makeTransformerEngineTensor(input_tensor);
 
-  NVTE_CHECK(amax.scalar_type() == torch_stable::ScalarType::Float,
-             "amax must be a float tensor");
+  NVTE_CHECK(amax.scalar_type() == torch_stable::ScalarType::Float, "amax must be a float tensor");
   NVTE_CHECK(amax.numel() == 1, "amax must have exactly one element");
   auto* amax_ptr = static_cast<float*>(amax.data_ptr());
   TensorWrapper fake_te_output(
@@ -28,10 +27,11 @@ void compute_amax(const torch_stable::Tensor& tensor, torch_stable::Tensor& amax
   nvte_compute_amax(te_input.data(), fake_te_output.data(), torch_stable::getCurrentCUDAStream());
 }
 
-void fused_amax_and_scale_update_after_reduction(
-    const torch_stable::Tensor& amax_reduction_buffer,
-    std::vector<torch_stable::Tensor> amax_histories, std::vector<torch_stable::Tensor> scales,
-    const std::string& amax_compute_algo, DType fp8_dtype, float margin) {
+void fused_amax_and_scale_update_after_reduction(const torch_stable::Tensor& amax_reduction_buffer,
+                                                 std::vector<torch_stable::Tensor> amax_histories,
+                                                 std::vector<torch_stable::Tensor> scales,
+                                                 const std::string& amax_compute_algo,
+                                                 DType fp8_dtype, float margin) {
   size_t num_tensors = amax_histories.size();
 
   // Allocate amax history and scale NVTETensors as batches
diff --git a/transformer_engine/pytorch/csrc/torch_stable.h b/transformer_engine/pytorch/csrc/torch_stable.h
index 19a15b7308..b8319051ec 100644
--- a/transformer_engine/pytorch/csrc/torch_stable.h
+++ b/transformer_engine/pytorch/csrc/torch_stable.h
@@ -44,10 +44,9 @@ inline Tensor contiguous(const Tensor &tensor) {
 
 inline cudaStream_t getCurrentCUDAStream() {
 #ifdef NVTE_WITH_TORCH_STABLE
-  return static_cast<cudaStream_t>(
-      torch::stable::accelerator::getCurrentStream(
-          torch::stable::accelerator::getCurrentDeviceIndex())
-          .nativeHandle());
+  return static_cast<cudaStream_t>(torch::stable::accelerator::getCurrentStream(
+                                       torch::stable::accelerator::getCurrentDeviceIndex())
+                                       .nativeHandle());
 #else
   return at::cuda::getCurrentCUDAStream();
 #endif