[PyTorch] Add torch_stable compatibility layer for incremental migration to the torch stable ABI

.gitignore

@@ -43,3 +43,4 @@ tensor_dumps/
 artifacts/
 .DS_Store
 .claude/
+transformer_engine/pytorch/_build_config.py

build_tools/pytorch.py

@@ -111,6 +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_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_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(
+        '"""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"
+    )
+
     # Construct PyTorch CUDA extension
     sources = [str(path) for path in sources]
     include_dirs = [str(path) for path in include_dirs]

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

transformer_engine/pytorch/csrc/common.cpp

@@ -62,6 +62,21 @@ 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();
+  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;
+}
+#endif
+
 std::unique_ptr<Quantizer> convert_quantizer(py::handle quantizer) {
   init_extension();
   if (quantizer.is_none()) {
@@ -165,6 +180,15 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor(at::Tensor tensor)
   return makeTransformerEngineTensor(tensor.data_ptr(), shape, dtype);
 }
 
+#ifdef NVTE_WITH_TORCH_STABLE
+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);
+}
+#endif
+
 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) {

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_stable.h"
 
 namespace transformer_engine::pytorch {
 
@@ -544,6 +546,10 @@ 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_stable::Tensor& tensor);
+#endif
+
 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 +587,10 @@ 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);
 
 // Flatten an N-D shape to 2D: {product(shape[:-1]), shape[-1]}.

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 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,
+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);
 

transformer_engine/pytorch/csrc/extensions/recipe.cpp

@@ -4,34 +4,32 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include <ATen/ATen.h>
-#include <ATen/cuda/CUDAContext.h>
-
 #include <string>
 
 #include "../extensions.h"
+#include "../torch_stable.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_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(
       /*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_stable::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,
+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();
@@ -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_stable::getCurrentCUDAStream());
 }
 
 }  // namespace transformer_engine::pytorch

transformer_engine/pytorch/csrc/extensions/stable_tensor_caster.h

@@ -0,0 +1,56 @@
+/*************************************************************************
+ * 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_
+
+/* 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 <pybind11/pybind11.h>
+
+#include "../torch_stable.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 || !torch::stable::is_tensor_pyobject(src.ptr())) {
+      return false;
+    }
+    value = torch::stable::tensor_from_pyobject(src.ptr());
+    return true;
+  }
+
+  static handle cast(const torch::stable::Tensor &src, return_value_policy, handle) {
+    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_

transformer_engine/pytorch/csrc/torch_stable.h

@@ -0,0 +1,57 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_
+#define TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_
+
+#include <cuda_runtime.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>
+#endif
+
+/* Compatibility layer for the incremental migration to the torch stable ABI.
+ * 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_stable {
+
+#ifdef NVTE_WITH_TORCH_STABLE
+using Tensor = torch::stable::Tensor;
+using ScalarType = torch::headeronly::ScalarType;
+#else
+using Tensor = at::Tensor;
+using ScalarType = at::ScalarType;
+#endif
+
+inline Tensor contiguous(const Tensor &tensor) {
+#ifdef NVTE_WITH_TORCH_STABLE
+  return torch::stable::contiguous(tensor);
+#else
+  return tensor.contiguous();
+#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_stable
+
+#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_TORCH_STABLE_H_