Port softmax ops to libtorch stable ABI

Proof of concept for migrating pybind11 functions to the PyTorch
stable ABI. Ports all 8 scaled softmax functions:

- Add stable_common.h with stable ABI helpers (tensor allocation,
  TensorWrapper construction, CUDA stream, dtype converters)
- Add registration.cpp with STABLE_TORCH_LIBRARY schema definitions
- Rewrite softmax.cpp: at::Tensor -> torch::stable::Tensor, use
  stable allocation and stream APIs, TORCH_BOX() for impl registration
- Remove softmax registrations from pybind.cpp
- Update Python callers to use torch.ops.transformer_engine_stable

The pattern is mechanical (API translation, no logic changes) and
establishes the template for porting the remaining ~70 Category A
functions that have no py::handle/py::object dependencies.

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

build_tools/pytorch.py

@@ -14,7 +14,15 @@
 
 def install_requirements() -> List[str]:
     """Install dependencies for TE/PyTorch extensions."""
-    return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic", "nvdlfw-inspect"]
+    return [
+        "torch>=2.10",
+        "einops",
+        "onnxscript",
+        "onnx",
+        "packaging",
+        "pydantic",
+        "nvdlfw-inspect",
+    ]
 
 
 def test_requirements() -> List[str]:

pyproject.toml

@@ -3,7 +3,7 @@
 # See LICENSE for license information.
 
 [build-system]
-requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1"]
+requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.10", "jax>=0.5.0", "flax>=0.7.1"]
 
 # Use legacy backend to import local packages in setup.py
 build-backend = "setuptools.build_meta:__legacy__"

transformer_engine/pytorch/__init__.py

@@ -13,7 +13,7 @@
 from transformer_engine.common import load_framework_extension
 from transformer_engine.pytorch.torch_version import torch_version
 
-assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}."
+assert torch_version() >= (2, 10), f"Minimum torch version 2.10 required. Found {torch_version()}."
 
 load_framework_extension("torch")
 from transformer_engine.pytorch.module import LayerNormLinear

transformer_engine/pytorch/attention/dot_product_attention/softmax.py

@@ -7,9 +7,10 @@
 from typing import Callable, Tuple, Union, Optional
 import torch
 from torch import nn
-import transformer_engine_torch as tex
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
+_ops = torch.ops.transformer_engine
+
 
 THREADS_PER_WARP = 32
 THREADS_PER_BLOCK = 128
@@ -47,7 +48,7 @@ class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
     def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor:
         """ScaledUpperTriangMaskedSoftmax fwd"""
         scale_t = torch.tensor([scale])
-        softmax_results = tex.scaled_upper_triang_masked_softmax_forward(inputs, scale_t[0])
+        softmax_results = _ops.scaled_upper_triang_masked_softmax_forward(inputs, scale_t[0])
 
         ctx.save_for_backward(softmax_results, scale_t)
         return softmax_results
@@ -56,7 +57,7 @@ def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor:
     def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
         """ScaledUpperTriangMaskedSoftmax bwd"""
         softmax_results, scale_t = ctx.saved_tensors
-        input_grads = tex.scaled_upper_triang_masked_softmax_backward(
+        input_grads = _ops.scaled_upper_triang_masked_softmax_backward(
             output_grads, softmax_results, scale_t[0]
         )
 
@@ -75,15 +76,15 @@ class ScaledAlignedCausalMaskedSoftmax(torch.autograd.Function):
     def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor:
         """ScaledAlignedCausalMaskedSoftmax fwd"""
         scale_t = torch.tensor([scale])
-        softmax_results = tex.scaled_aligned_causal_masked_softmax_forward(inputs, scale_t[0])
+        softmax_results = _ops.scaled_aligned_causal_masked_softmax_forward(inputs, scale_t[0])
         ctx.save_for_backward(softmax_results, scale_t)
         return softmax_results
 
     @staticmethod
     def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]:
         """ScaledAlignedCausalMaskedSoftmax bwd"""
         softmax_results, scale_t = ctx.saved_tensors
-        input_grads = tex.scaled_aligned_causal_masked_softmax_backward(
+        input_grads = _ops.scaled_aligned_causal_masked_softmax_backward(
             output_grads, softmax_results, scale_t[0]
         )
 
@@ -103,7 +104,7 @@ def forward(ctx, inputs: torch.Tensor, mask: torch.Tensor, scale: float) -> torc
         """ScaledMaskedSoftmax fwd"""
         scale_t = torch.tensor([scale])
 
-        softmax_results = tex.scaled_masked_softmax_forward(inputs, mask, scale_t[0])
+        softmax_results = _ops.scaled_masked_softmax_forward(inputs, mask, scale_t[0])
         ctx.save_for_backward(softmax_results, scale_t)
         return softmax_results
 
@@ -112,7 +113,7 @@ def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None]
         """ScaledMaskedSoftmax bwd"""
         softmax_results, scale_t = ctx.saved_tensors
 
-        input_grads = tex.scaled_masked_softmax_backward(output_grads, softmax_results, scale_t[0])
+        input_grads = _ops.scaled_masked_softmax_backward(output_grads, softmax_results, scale_t[0])
         return input_grads, None, None
 
 
@@ -128,7 +129,7 @@ def forward(ctx, inputs: torch.Tensor, scale: float) -> torch.Tensor:
         """ScaledSoftmax fwd"""
         scale_t = torch.tensor([scale])
 
-        softmax_results = tex.scaled_softmax_forward(inputs, scale_t[0])
+        softmax_results = _ops.scaled_softmax_forward(inputs, scale_t[0])
         ctx.save_for_backward(softmax_results, scale_t)
         return softmax_results
 
@@ -137,7 +138,7 @@ def backward(ctx, output_grads: torch.Tensor) -> Tuple[Union[torch.Tensor, None]
         """ScaledSoftmax bwd"""
         softmax_results, scale_t = ctx.saved_tensors
 
-        input_grads = tex.scaled_softmax_backward(output_grads, softmax_results, scale_t[0])
+        input_grads = _ops.scaled_softmax_backward(output_grads, softmax_results, scale_t[0])
         return input_grads, None, None
 
 

transformer_engine/pytorch/csrc/extensions.h

@@ -349,32 +349,6 @@ py::object dropout_bwd(const at::Tensor &grad_output, const at::Tensor &mask,
                        const float dropout_probability,
                        std::optional<at::Tensor> grad_input = std::nullopt);
 
-/***************************************************************************************************
- * Softmax
- **************************************************************************************************/
-
-at::Tensor scaled_softmax_forward(at::Tensor input, float scale_factor);
-
-at::Tensor scaled_softmax_backward(at::Tensor output_grad_, at::Tensor softmax_results_,
-                                   float scale_factor);
-
-at::Tensor scaled_masked_softmax_forward(at::Tensor input, at::Tensor mask, float scale_factor);
-
-at::Tensor scaled_masked_softmax_backward(at::Tensor output_grad_, at::Tensor softmax_results_,
-                                          float scale_factor);
-
-at::Tensor scaled_upper_triang_masked_softmax_forward(at::Tensor input, float scale_factor);
-
-at::Tensor scaled_upper_triang_masked_softmax_backward(at::Tensor output_grads_,
-                                                       at::Tensor softmax_results_,
-                                                       float scale_factor);
-
-at::Tensor scaled_aligned_causal_masked_softmax_forward(at::Tensor input, float scale_factor);
-
-at::Tensor scaled_aligned_causal_masked_softmax_backward(at::Tensor output_grads_,
-                                                         at::Tensor softmax_results_,
-                                                         float scale_factor);
-
 /***************************************************************************************************
  * FP8 recipe
  **************************************************************************************************/

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -232,32 +232,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("moe_unpermute_bwd", transformer_engine::pytorch::moe_unpermute_bwd, "MOE unpermute BWD",
         py::call_guard<py::gil_scoped_release>());
 
-  // Softmax functions
-  m.def("scaled_softmax_forward", &transformer_engine::pytorch::scaled_softmax_forward,
-        "Scaled Softmax FWD", py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_softmax_backward", &transformer_engine::pytorch::scaled_softmax_backward,
-        "Scaled Softmax BWD", py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_masked_softmax_forward",
-        &transformer_engine::pytorch::scaled_masked_softmax_forward, "Scaled Masked Softmax FWD",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_masked_softmax_backward",
-        &transformer_engine::pytorch::scaled_masked_softmax_backward, "Scaled Masked Softmax BWD",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_upper_triang_masked_softmax_forward",
-        &transformer_engine::pytorch::scaled_upper_triang_masked_softmax_forward,
-        "Scaled Upper-Triangular Masked Softmax FWD", py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_upper_triang_masked_softmax_backward",
-        &transformer_engine::pytorch::scaled_upper_triang_masked_softmax_backward,
-        "Scaled Upper-Triangular Masked Softmax BWD", py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_aligned_causal_masked_softmax_forward",
-        &transformer_engine::pytorch::scaled_aligned_causal_masked_softmax_forward,
-        "Scaled Bottom-Right Corner Aligned Masked Softmax FWD",
-        py::call_guard<py::gil_scoped_release>());
-  m.def("scaled_aligned_causal_masked_softmax_backward",
-        &transformer_engine::pytorch::scaled_aligned_causal_masked_softmax_backward,
-        "Scaled Bottom-Right Corner Aligned Masked Softmax BWD",
-        py::call_guard<py::gil_scoped_release>());
-
   // Other granular functions
   m.def("layernorm_fwd", &transformer_engine::pytorch::layernorm_fwd, "LayerNorm", py::arg("input"),
         py::arg("weight"), py::arg("bias"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"),

transformer_engine/pytorch/csrc/extensions/registration.cpp

@@ -0,0 +1,30 @@
+/*************************************************************************
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ *
+ * See LICENSE for license information.
+ ************************************************************************/
+
+#include "../stable_common.h"
+
+// This file defines the transformer_engine library namespace.
+// All other stable ABI files use STABLE_TORCH_LIBRARY_FRAGMENT to add schemas
+// and STABLE_TORCH_LIBRARY_IMPL to add implementations.
+STABLE_TORCH_LIBRARY(transformer_engine, m) {
+  // Softmax ops
+  m.def("scaled_softmax_forward(Tensor input, float scale_factor) -> Tensor");
+  m.def(
+      "scaled_softmax_backward(Tensor output_grad, Tensor softmax_results, float scale_factor) -> "
+      "Tensor");
+  m.def("scaled_masked_softmax_forward(Tensor input, Tensor mask, float scale_factor) -> Tensor");
+  m.def(
+      "scaled_masked_softmax_backward(Tensor output_grad, Tensor softmax_results, float "
+      "scale_factor) -> Tensor");
+  m.def("scaled_upper_triang_masked_softmax_forward(Tensor input, float scale_factor) -> Tensor");
+  m.def(
+      "scaled_upper_triang_masked_softmax_backward(Tensor output_grads, Tensor softmax_results, "
+      "float scale_factor) -> Tensor");
+  m.def("scaled_aligned_causal_masked_softmax_forward(Tensor input, float scale_factor) -> Tensor");
+  m.def(
+      "scaled_aligned_causal_masked_softmax_backward(Tensor output_grad, Tensor softmax_results, "
+      "float scale_factor) -> Tensor");
+}