/*************************************************************************
 * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include "pybind.h"

#include <pybind11/cast.h>
#include <pybind11/detail/common.h>
#include <pybind11/functional.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <memory>
#include <optional>
#include <vector>

#include "../common.h"
#include "../extensions.h"
#include "common.h"

namespace transformer_engine::pytorch {

PyTypeObject *Float8TensorPythonClass = nullptr;  /// TODO Remove
PyTypeObject *Float8TensorStoragePythonClass = nullptr;
PyTypeObject *Float8QuantizerClass = nullptr;
PyTypeObject *Float8CurrentScalingQuantizerClass = nullptr;
PyTypeObject *MXFP8TensorPythonClass = nullptr;  /// TODO Remove
PyTypeObject *MXFP8TensorStoragePythonClass = nullptr;
PyTypeObject *MXFP8QuantizerClass = nullptr;
PyTypeObject *Float8BlockwiseQTensorPythonClass = nullptr;
PyTypeObject *Float8BlockwiseQTensorStoragePythonClass = nullptr;
PyTypeObject *Float8BlockwiseQuantizerClass = nullptr;
PyTypeObject *NVFP4TensorPythonClass = nullptr;
PyTypeObject *NVFP4TensorStoragePythonClass = nullptr;
PyTypeObject *NVFP4QuantizerClass = nullptr;
PyTypeObject *GroupedTensorPythonClass = nullptr;
PyTypeObject *GroupedTensorStoragePythonClass = nullptr;
std::once_flag extension_init_flag;

void init_float8_extension() {
  auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.float8_tensor");
  Float8QuantizerClass =
      reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "Float8Quantizer"));
  Float8CurrentScalingQuantizerClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_module.ptr(), "Float8CurrentScalingQuantizer"));
  Float8TensorPythonClass =
      reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "Float8Tensor"));
  auto fp8_base_module =
      py::module_::import("transformer_engine.pytorch.tensor.storage.float8_tensor_storage");
  Float8TensorStoragePythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_base_module.ptr(), "Float8TensorStorage"));
  NVTE_CHECK(Float8TensorPythonClass != nullptr,
             "Internal error: could not initialize pyTorch Float8 extension.");
}

void init_mxfp8_extension() {
  auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.mxfp8_tensor");
  MXFP8QuantizerClass =
      reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "MXFP8Quantizer"));
  MXFP8TensorPythonClass =
      reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(fp8_module.ptr(), "MXFP8Tensor"));
  auto fp8_base_module =
      py::module_::import("transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage");
  MXFP8TensorStoragePythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_base_module.ptr(), "MXFP8TensorStorage"));
  NVTE_CHECK(MXFP8TensorPythonClass != nullptr,
             "Internal error: could not initialize pyTorch MXFP8 extension.");
}

void init_float8blockwise_extension() {
  auto fp8_module =
      py::module_::import("transformer_engine.pytorch.tensor.float8_blockwise_tensor");
  auto fp8_base_module = py::module_::import(
      "transformer_engine.pytorch.tensor.storage.float8_blockwise_tensor_storage");
  Float8BlockwiseQuantizerClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_module.ptr(), "Float8BlockQuantizer"));
  Float8BlockwiseQTensorStoragePythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_base_module.ptr(), "Float8BlockwiseQTensorStorage"));
  Float8BlockwiseQTensorPythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(fp8_module.ptr(), "Float8BlockwiseQTensor"));

  NVTE_CHECK(Float8BlockwiseQuantizerClass != nullptr,
             "Internal error: could not initialize pyTorch float8blockwise extension.");
  NVTE_CHECK(Float8BlockwiseQTensorStoragePythonClass != nullptr,
             "Internal error: could not initialize pyTorch float8blockwise extension.");
  NVTE_CHECK(Float8BlockwiseQTensorPythonClass != nullptr,
             "Internal error: could not initialize pyTorch float8blockwise extension.");
}

void init_nvfp4_extensions() {
  auto nvfp4_module = py::module_::import("transformer_engine.pytorch.tensor.nvfp4_tensor");
  NVFP4QuantizerClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(nvfp4_module.ptr(), "NVFP4Quantizer"));
  NVFP4TensorPythonClass =
      reinterpret_cast<PyTypeObject *>(PyObject_GetAttrString(nvfp4_module.ptr(), "NVFP4Tensor"));
  auto nvfp4_base_module =
      py::module_::import("transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage");
  NVFP4TensorStoragePythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(nvfp4_base_module.ptr(), "NVFP4TensorStorage"));
  NVTE_CHECK(NVFP4TensorPythonClass != nullptr,
             "Internal error: could not initialize pyTorch NVFP4 extension.");
}

void init_grouped_tensor_extension() {
  if (GroupedTensorPythonClass && GroupedTensorStoragePythonClass) return;
  auto grouped_tensor_module =
      py::module_::import("transformer_engine.pytorch.tensor.grouped_tensor");
  GroupedTensorPythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(grouped_tensor_module.ptr(), "GroupedTensor"));
  auto grouped_tensor_storage_module =
      py::module_::import("transformer_engine.pytorch.tensor.storage.grouped_tensor_storage");
  GroupedTensorStoragePythonClass = reinterpret_cast<PyTypeObject *>(
      PyObject_GetAttrString(grouped_tensor_storage_module.ptr(), "GroupedTensorStorage"));
  NVTE_CHECK(GroupedTensorPythonClass != nullptr,
             "Internal error: could not initialize pyTorch grouped tensor extension.");
  NVTE_CHECK(GroupedTensorStoragePythonClass != nullptr,
             "Internal error: could not initialize pyTorch grouped tensor extension.");
}

void init_extension() {
  pybind11::gil_scoped_acquire gil;
  std::call_once(extension_init_flag, []() {
    init_float8_extension();
    init_mxfp8_extension();
    init_float8blockwise_extension();
    init_nvfp4_extensions();
    init_grouped_tensor_extension();
  });
}

// Pybind11 registrations for the fused MoE router kernels. Split out of
// PYBIND11_MODULE() to keep that function under the cpplint readability/fn_size
// limit.
void init_router_bindings(pybind11::module &m) {
  pybind11::enum_<NVTERoutingMapFormat>(m, "NVTERoutingMapFormat", pybind11::module_local())
      .value("BYTEMAP", NVTE_ROUTING_MAP_FORMAT_BYTEMAP)
      .value("BITMAP_U8", NVTE_ROUTING_MAP_FORMAT_BITMAP_U8);
  m.def("fused_topk_with_score_function_fwd", &fused_topk_with_score_function_fwd,
        py::arg("logits"), py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"),
        py::arg("group_topk"), py::arg("scaling_factor"), py::arg("score_function"),
        py::arg("expert_bias"),
        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
        py::arg("topk_indices") = std::nullopt, "Fused topk with score function fwd");
  m.def("fused_topk_with_score_function_bwd", &fused_topk_with_score_function_bwd,
        py::arg("routing_map"), py::arg("intermediate_output"), py::arg("grad_probs"),
        py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"),
        py::arg("scaling_factor"), py::arg("score_function"), py::arg("use_dense_indices") = false,
        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
        "Fused topk with score function bwd");
  m.def("fused_score_for_moe_aux_loss_fwd", &fused_score_for_moe_aux_loss_fwd, py::arg("logits"),
        py::arg("topk"), py::arg("score_function"),
        py::arg("routing_map_format") = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
        "Fused aux loss with score function fwd");
  m.def("fused_score_for_moe_aux_loss_bwd", &fused_score_for_moe_aux_loss_bwd,
        py::arg("intermediate_output"), py::arg("grad_scores"), py::arg("grad_logits"),
        py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function bwd");
  m.def("fused_moe_aux_loss_fwd", &fused_moe_aux_loss_fwd, py::arg("probs"),
        py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"),
        py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"),
        "Fused aux loss fwd (host-int total_num_tokens, host-folded C_coeff)");
  m.def("fused_moe_aux_loss_fwd_graph_safe", &fused_moe_aux_loss_fwd_graph_safe, py::arg("probs"),
        py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"),
        py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), py::arg("coeff"),
        "Fused aux loss fwd (device-tensor total_num_tokens, CUDA-graph-safe)");
  m.def("fused_moe_aux_loss_bwd", &fused_moe_aux_loss_bwd, py::arg("Const_buf"),
        py::arg("tokens_per_expert"), py::arg("num_rows"), py::arg("num_cols"),
        py::arg("grad_aux_loss"), "Fused aux loss bwd");
}

void bind_quantize_with_amax_extensions(py::module_ &m) {
  m.def("nvfp4_quantize_with_amax", nvfp4_quantize_with_amax, py::arg("tensor"),
        py::arg("quantizer"), py::arg("rowwise_amax"), py::arg("columnwise_amax"));
  m.def("nvfp4_group_quantize_with_amax", nvfp4_group_quantize_with_amax, py::arg("tensor"),
        py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
        py::arg("last_dims") = py::none(), py::arg("rowwise_amax"), py::arg("columnwise_amax"),
        py::arg("tensor_offsets") = py::none());
}

}  // namespace transformer_engine::pytorch

#include "common/util/pybind_helper.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m)

  // Register __eq__/__ne__ on the pybind ``DType`` enum so it compares by integer.
  py::object dtype_class = m.attr("DType");
  dtype_class.attr("__eq__") = py::cpp_function(
      [](transformer_engine::DType self, py::object other) -> py::object {
        return py::cast(static_cast<int>(self) == other.cast<int>());
      },
      py::is_method(dtype_class));
  dtype_class.attr("__ne__") = py::cpp_function(
      [](transformer_engine::DType self, py::object other) -> py::object {
        return py::cast(static_cast<int>(self) != other.cast<int>());
      },
      py::is_method(dtype_class));

  m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"),
        py::arg("output") = py::none(), py::arg("noop") = py::none());
  m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"),
        py::arg("otype"));
  m.def("create_empty_quantized_tensor",
        &transformer_engine::pytorch::create_empty_quantized_tensor,
        "Create an empty quantized tensor", py::arg("quantizer"), py::arg("shape"),
        py::arg("dtype"), py::arg("device"), py::arg("pin_memory"));
  m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"),
        py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
        py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(),
        py::arg("noop_flag") = py::none());
  transformer_engine::pytorch::bind_quantize_with_amax_extensions(m);
  m.def("group_dequantize", transformer_engine::pytorch::group_dequantize,
        "Dequantize group tensor", py::arg("input"), py::arg("otype"));
  m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
        py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
        py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
  m.def("group_requantize_inplace", transformer_engine::pytorch::group_requantize_inplace,
        "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
        "its rowwise scales for GEMM, in place",
        py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
        py::arg("otype"), py::arg("tensor_offsets") = py::none(),
        py::arg("return_dequantized") = false);
  m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
        "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer"));
  m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)",
        py::arg("A"), py::arg("transA"), py::arg("B"), py::arg("transB"), py::arg("D"),
        py::arg("quantizer"), py::arg("output_dtype"), py::arg("bias"), py::arg("bias_type"),
        py::arg("gelu"), py::arg("gelu_in"), py::arg("grad"), py::arg("workspace"),
        py::arg("workspace_size"), py::arg("accumulate"), py::arg("use_split_accumulator"),
        py::arg("comm_overlap") = nullptr, py::arg("comm_type") = std::nullopt,
        py::arg("extra_output") = std::nullopt, py::arg("bulk_overlap") = false,
        py::arg("alpha") = 1.0f, py::arg("beta") = std::nullopt);
  /* GLU (sigmoid gate) */
  m.def("glu", transformer_engine::pytorch::glu, "GLU activation", py::arg("input"),
        py::arg("quantizer"));
  /* GELU and variants*/
  m.def("gelu", transformer_engine::pytorch::gelu, "GeLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("geglu", transformer_engine::pytorch::geglu, "GeGLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("qgelu", transformer_engine::pytorch::qgelu, "QuickGELU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("qgeglu", transformer_engine::pytorch::qgeglu, "QuickGeGLU activation", py::arg("input"),
        py::arg("quantizer"));
  /* ReLU and variants */
  m.def("relu", transformer_engine::pytorch::relu, "ReLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("reglu", transformer_engine::pytorch::reglu, "ReGLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("srelu", transformer_engine::pytorch::srelu, "Squared ReLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("sreglu", transformer_engine::pytorch::sreglu, "Squared ReGLU activation", py::arg("input"),
        py::arg("quantizer"));
  /* SwiGLU and variants */
  m.def("silu", transformer_engine::pytorch::silu, "SiLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("swiglu", transformer_engine::pytorch::swiglu, "SwiGLU activation", py::arg("input"),
        py::arg("quantizer"));
  m.def("clamped_swiglu", transformer_engine::pytorch::clamped_swiglu,
        "SwiGLU activation used in GPT OSS", py::arg("input"), py::arg("quantizer"),
        py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f);
  /* Backward of GLU */
  m.def("dglu", transformer_engine::pytorch::dglu, "Backward of GLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  /* Backward of GELU and variants */
  m.def("dgelu", transformer_engine::pytorch::dgelu, "Backward of GeLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dgeglu", transformer_engine::pytorch::dgeglu, "Backward of GeGLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dqgelu", transformer_engine::pytorch::dqgelu, "Backward of QuickGELU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dqgeglu", transformer_engine::pytorch::dqgeglu, "Backward of QuickGeGLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  /* Backward of ReLU and variants */
  m.def("drelu", transformer_engine::pytorch::drelu, "Backward of ReLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dreglu", transformer_engine::pytorch::dreglu, "Backward of ReGLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dsrelu", transformer_engine::pytorch::dsrelu, "Backward of Squared ReLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dsreglu", transformer_engine::pytorch::dsreglu, "Backward of Squared ReGLU",
        py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
  /* Backward of SiLU and variants */
  m.def("dsilu", transformer_engine::pytorch::dsilu, "Backward of SiLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dswiglu", transformer_engine::pytorch::dswiglu, "Backward of SwiGLU", py::arg("grad"),
        py::arg("fwd_input"), py::arg("quantizer"));
  m.def("clamped_dswiglu", transformer_engine::pytorch::clamped_dswiglu,
        "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"),
        py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
        py::arg("glu_linear_offset") = 1.0f);
  /* Scaled activation */
  m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation",
        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"),
        py::arg("glu_interleave_size") = 0);
  m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu,
        "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"),
        py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
        py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0);
  m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation",
        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"));
  m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward",
        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
  m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu,
        "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"),
        py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f,
        py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f,
        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
  m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward",
        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
        py::arg("compute_scale_grad") = true);
  /* DBias + DAct fusions*/
  m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize",
        py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dbias_dsilu", transformer_engine::pytorch::dbias_dsilu, "DSiLU + DBias + Quantize",
        py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dbias_drelu", transformer_engine::pytorch::dbias_drelu, "DReLU + DBias + Quantize",
        py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dbias_dqgelu", transformer_engine::pytorch::dbias_dqgelu, "DQGeLU + DBias + Quantize",
        py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
  m.def("dbias_dsrelu", transformer_engine::pytorch::dbias_dsrelu,
        "DSquaredReLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"),
        py::arg("quantizer"));

#ifdef NVTE_WITH_NCCL_EP
  transformer_engine::pytorch::register_ep_bindings(m);
#endif  // NVTE_WITH_NCCL_EP

  // Permutation functions
  m.def("moe_permute_fwd", transformer_engine::pytorch::moe_permute_fwd, "MOE permute FWD",
        py::call_guard<py::gil_scoped_release>());
  m.def("moe_permute_bwd", transformer_engine::pytorch::moe_permute_bwd, "MOE permute BWD",
        py::call_guard<py::gil_scoped_release>());
  m.def("moe_unpermute_fwd", transformer_engine::pytorch::moe_unpermute_fwd, "MOE unpermute FWD",
        py::call_guard<py::gil_scoped_release>());
  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"),
        py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma"));
  m.def("layernorm_bwd", &transformer_engine::pytorch::layernorm_bwd, "Backward of LayerNorm");
  m.def("rmsnorm_fwd", &transformer_engine::pytorch::rmsnorm_fwd, "RMSNorm", py::arg("input"),
        py::arg("weight"), py::arg("eps"), py::arg("ln_out"), py::arg("quantizer"),
        py::arg("otype"), py::arg("sm_margin"), py::arg("zero_centered_gamma"));
  m.def("rmsnorm_bwd", &transformer_engine::pytorch::rmsnorm_bwd, "Backward of RMSNorm");
  m.def("rmsnorm_bwd_add", &transformer_engine::pytorch::rmsnorm_bwd_add,
        "Fused backward of RMSNorm + add");
  m.def("multi_tensor_quantize", &transformer_engine::pytorch::multi_tensor_quantize,
        "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list"));
  m.def("split_quantize", &transformer_engine::pytorch::split_quantize,
        "Split and multi-tensor quantize", py::arg("tensor"), py::arg("split_sections"),
        py::arg("quantizer_list"), py::arg("disable_bulk_allocation") = false);
  m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size,
        "Required workspace size for grouped GEMM setup");
  m.def("te_general_grouped_gemm", &transformer_engine::pytorch::te_general_grouped_gemm,
        "Grouped GEMM");
  m.def("te_general_grouped_gemm_for_grouped_tensor",
        &transformer_engine::pytorch::te_general_grouped_gemm_for_grouped_tensor,
        "Grouped GEMM for GroupedTensor");
  m.def("te_general_grouped_gemm_for_discrete_in",
        &transformer_engine::pytorch::te_general_grouped_gemm_for_discrete_in,
        "Grouped GEMM for discrete A input list");
  m.def("te_general_grouped_gemm_for_discrete_out",
        &transformer_engine::pytorch::te_general_grouped_gemm_for_discrete_out,
        "Grouped GEMM for discrete output list");
  m.def("fp8_transpose", &transformer_engine::pytorch::fp8_transpose, "Transpose with FP8 I/O",
        py::arg("input"), py::arg("dtype"), py::kw_only(), py::arg("out"),
        py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_data_transpose", &transformer_engine::pytorch::nvfp4_data_transpose,
        "Transpose NVFP4 packed data with nibble repacking", py::arg("input"), py::kw_only(),
        py::arg("out"), py::call_guard<py::gil_scoped_release>());
  m.def(
      "nvfp4_2d_scale_transpose", &transformer_engine::pytorch::nvfp4_2d_scale_transpose,
      "Transpose NVFP4 tile-level scales (E4M3 stored as uint8) from rowwise to columnwise format",
      py::arg("input"), py::arg("output"), py::arg("M_tiles"), py::arg("K_tiles"),
      py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_expand_scale_to_fp8", &transformer_engine::pytorch::nvfp4_expand_scale_to_fp8,
        "Expand tile-level scales to row-level scales and convert to FP8 E4M3", py::arg("input"),
        py::arg("output"), py::arg("tile_rows"), py::arg("tile_cols"), py::arg("rows_padded"),
        py::arg("block_len"), py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_compute_per_block_scale",
        &transformer_engine::pytorch::nvfp4_compute_per_block_scale,
        "Compute per-block decode scale from block amax and global amax", py::arg("block_amax"),
        py::arg("scale"), py::arg("global_amax"), py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_compute_global_scale", &transformer_engine::pytorch::nvfp4_compute_global_scale,
        "Compute global encode scale from global amax", py::arg("global_amax"),
        py::arg("global_scale"), py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_fused_scale", &transformer_engine::pytorch::nvfp4_fused_scale,
        "Fused kernel: compute per-block decode scale, copy global amax, expand to row-level FP8",
        py::arg("block_amax"), py::arg("global_amax"), py::arg("per_block_scale"),
        py::arg("target_scale"), py::arg("target_amax"), py::arg("tile_rows"), py::arg("tile_cols"),
        py::arg("rows_padded"), py::arg("block_len"), py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_multi_tensor_fused_scale",
        &transformer_engine::pytorch::nvfp4_multi_tensor_fused_scale,
        "Batched fused scale: compute per-block decode scale, copy global amax, expand to FP8 for "
        "multiple tensors",
        py::arg("block_amax_list"), py::arg("global_amax_list"), py::arg("per_block_scale_list"),
        py::arg("target_scale_list"), py::arg("target_amax_list"), py::arg("tile_rows_list"),
        py::arg("tile_cols_list"), py::arg("rows_padded_list"), py::arg("block_len"),
        py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_2d_multi_tensor_transpose",
        &transformer_engine::pytorch::nvfp4_2d_multi_tensor_transpose,
        "Batched NVFP4 columnwise creation: transpose data and scales for multiple tensors",
        py::arg("rowwise_data_list"), py::arg("columnwise_data_list"),
        py::arg("rowwise_scale_inv_list"), py::arg("columnwise_scale_inv_list"), py::arg("M_list"),
        py::arg("K_list"), py::call_guard<py::gil_scoped_release>());
  m.def("swap_first_dims", &transformer_engine::pytorch::swap_first_dims,
        "Swap first two tensor dimensions", py::arg("tensor"), py::kw_only(), py::arg("out"),
        py::call_guard<py::gil_scoped_release>());
  m.def("get_fused_attn_backend", &transformer_engine::pytorch::get_fused_attn_backend,
        "Get Fused Attention backend", py::call_guard<py::gil_scoped_release>());
  m.def("compute_amax", &transformer_engine::pytorch::compute_amax,
        "Compute absolute max value in tensor", py::arg("input"), py::arg("amax"),
        py::call_guard<py::gil_scoped_release>());
  m.def("fused_amax_and_scale_update_after_reduction",
        &transformer_engine::pytorch::fused_amax_and_scale_update_after_reduction,
        "Update amax history and FP8 scale/scale_inv after reduction",
        py::call_guard<py::gil_scoped_release>());
  m.def("fp8_block_scaling_compute_partial_amax",
        &transformer_engine::pytorch::fp8_block_scaling_compute_partial_amax,
        "Compute partial amax from master weights for fp8 block scaling", py::arg("tensor"),
        py::arg("amax"), py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len"),
        py::call_guard<py::gil_scoped_release>());
  m.def("fp8_block_scaling_partial_cast",
        &transformer_engine::pytorch::fp8_block_scaling_partial_cast,
        "Partial cast from master weights for fp8 block scaling", py::arg("inp"), py::arg("out"),
        py::arg("scale"), py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len"),
        py::arg("out_dtype"), py::call_guard<py::gil_scoped_release>());

  // NVFP4 2D
  m.def("nvfp4_2d_compute_partial_amax",
        &transformer_engine::pytorch::nvfp4_2d_compute_partial_amax,
        "Compute partial amax from master weights for NVFP4 2D", py::arg("tensor"), py::arg("amax"),
        py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len") = 16,
        py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_multi_tensor_compute_partial_amax",
        &transformer_engine::pytorch::nvfp4_multi_tensor_compute_partial_amax,
        "Batched compute partial and global amax from master weights for NVFP4 2D",
        py::arg("master_weight_list"), py::arg("partial_amax_list"), py::arg("global_amax_list"),
        py::arg("h_list"), py::arg("w_list"), py::arg("start_offset_list"),
        py::arg("block_len") = 16, py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_2d_partial_cast", &transformer_engine::pytorch::nvfp4_2d_partial_cast,
        "Partial cast from master weights for NVFP4 2D", py::arg("inp"), py::arg("out"),
        py::arg("scale"), py::arg("global_scale"), py::arg("h"), py::arg("w"),
        py::arg("start_offset"), py::arg("block_len") = 16,
        py::call_guard<py::gil_scoped_release>());
  m.def("nvfp4_multi_tensor_2d_partial_cast",
        &transformer_engine::pytorch::nvfp4_multi_tensor_2d_partial_cast,
        "Batched partial cast from master weights for NVFP4 2D", py::arg("inp_list"),
        py::arg("out_list"), py::arg("scale_list"), py::arg("global_scale_list"), py::arg("h_list"),
        py::arg("w_list"), py::arg("start_offset_list"), py::arg("block_len") = 16,
        py::call_guard<py::gil_scoped_release>());
  m.def("mxfp8_scaling_compute_partial_amax",
        &transformer_engine::pytorch::mxfp8_scaling_compute_partial_amax,
        "Compute partial amax from master weights for fp8 mxfp8 scaling", py::arg("input"),
        py::arg("amax_rowwise"), py::arg("amax_colwise"), py::arg("rows"), py::arg("cols"),
        py::arg("start_offset"), py::call_guard<py::gil_scoped_release>());
  m.def("mxfp8_scaling_partial_cast", &transformer_engine::pytorch::mxfp8_scaling_partial_cast,
        "Partial cast from master weights for fp8 mxfp8 scaling", py::arg("input"),
        py::arg("output_rowwise"), py::arg("output_colwise"), py::arg("scale_inv_rowwise"),
        py::arg("scale_inv_colwise"), py::arg("rows"), py::arg("cols"), py::arg("start_offset"),
        py::call_guard<py::gil_scoped_release>());
  m.def("fused_multi_row_padding", &transformer_engine::pytorch::fused_multi_row_padding,
        "Fused Multi-tensor padding", py::call_guard<py::gil_scoped_release>());
  m.def("fused_multi_row_unpadding", &transformer_engine::pytorch::fused_multi_row_unpadding,
        "Fused Multi-tensor unpadding", py::call_guard<py::gil_scoped_release>());
  m.def("swizzle_scales_for_gemm_", &transformer_engine::pytorch::inplace_swizzle_scale_for_gemm,
        "Convert tensor block scales into GEMM swizzled format");
  m.def("multi_tensor_swizzle_scales_for_gemm_",
        &transformer_engine::pytorch::inplace_multi_tensor_swizzle_scales_for_gemm,
        "Convert multiple tensors' block scales into GEMM swizzled format", py::arg("tensors"),
        py::arg("rowwise_usage"), py::arg("columnwise_usage"));
  m.def(
      "multi_tensor_swizzle_scales_for_gemm_unchecked_",
      &transformer_engine::pytorch::inplace_multi_tensor_swizzle_scales_for_gemm_unchecked,
      "Convert multiple tensors' block scales into GEMM swizzled format (skip scale shape checks)",
      py::arg("tensors"), py::arg("rowwise_usage"), py::arg("columnwise_usage"));
  m.def("grouped_swizzle_for_gemm", &transformer_engine::pytorch::grouped_swizzle_for_gemm,
        "In-place swizzle of grouped tensor scales for GEMM", py::arg("tensor"), py::arg("rowwise"),
        py::arg("columnwise"));

  // Tensor allocation
  m.def("bulk_allocate", &transformer_engine::pytorch::bulk_allocate,
        "Allocate tensors backed by a single contiguous buffer", py::arg("shapes"),
        py::arg("dtypes"), py::arg("device") = py::none(), py::arg("alignments") = py::none(),
        py::call_guard<py::gil_scoped_release>());

  // attention kernels
  m.def("fa_prepare_fwd", &transformer_engine::pytorch::fa_prepare_fwd,
        "Prepare QKV for Flash Attention", py::call_guard<py::gil_scoped_release>());
  m.def("fa_prepare_bwd", &transformer_engine::pytorch::fa_prepare_bwd,
        "Backward of QKV preparation for Flash Attention",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_transpose_to_bhsd",
        &transformer_engine::pytorch::multi_tensor_transpose_to_bhsd,
        "Permute multiple tensors from BSHD/SBHD to BHSD.", py::arg("inputs"),
        py::arg("original_format"), py::arg("outputs") = std::vector<std::optional<at::Tensor>>{},
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_pad_last_dim", &transformer_engine::pytorch::multi_tensor_pad_last_dim,
        "Pad multiple tensors' last dimension to a common alignment.", py::arg("inputs"),
        py::arg("alignment"), py::call_guard<py::gil_scoped_release>());
  m.def("fused_attn_fwd", &transformer_engine::pytorch::fused_attn_fwd,
        "Fused Attention FP8/BF16/FP16 FWD with separate Q, K and V");
  m.def("fused_attn_bwd", &transformer_engine::pytorch::fused_attn_bwd,
        "Fused Attention FP8/BF16/FP16 BWD with separate Q, K and V");
  m.def("copy_to_kv_cache", &transformer_engine::pytorch::copy_to_kv_cache,
        "Copy new KV tokens to KV cache", py::call_guard<py::gil_scoped_release>());
  m.def("convert_thd_to_bshd", &transformer_engine::pytorch::convert_thd_to_bshd,
        "Convert a tensor from THD to BSHD", py::call_guard<py::gil_scoped_release>());
  m.def("convert_bshd_to_thd", &transformer_engine::pytorch::convert_bshd_to_thd,
        "Convert a tesnor from BSHD to THD", py::call_guard<py::gil_scoped_release>());

  // fused apply rope
  m.def("fused_rope_forward", &transformer_engine::pytorch::fused_rope_forward,
        "Fused Apply RoPE FWD", py::call_guard<py::gil_scoped_release>());
  m.def("fused_rope_backward", &transformer_engine::pytorch::fused_rope_backward,
        "Fused Apply RoPE BWD", py::call_guard<py::gil_scoped_release>());
  m.def("fused_qkv_rope_forward", &transformer_engine::pytorch::fused_qkv_rope_forward,
        "Fused Apply QKV RoPE FWD", py::call_guard<py::gil_scoped_release>());
  m.def("fused_qkv_rope_backward", &transformer_engine::pytorch::fused_qkv_rope_backward,
        "Fused Apply QKV RoPE BWD", py::call_guard<py::gil_scoped_release>());

  // fused router
  transformer_engine::pytorch::init_router_bindings(m);

  // Dropout
  m.def("dropout_fwd", transformer_engine::pytorch::dropout_fwd, "Dropout forward with 8-bit RNG",
        py::arg("input"), py::arg("dropout_probability"), py::arg("out") = std::nullopt);
  m.def("dropout_bwd", transformer_engine::pytorch::dropout_bwd, "Dropout backward with 8-bit RNG",
        py::arg("grad_output"), py::arg("mask"), py::arg("dropout_probability"),
        py::arg("grad_input") = std::nullopt);

  // Misc
  m.def("get_cublasLt_version", &transformer_engine::pytorch::get_cublasLt_version,
        "Get cublasLt version", py::call_guard<py::gil_scoped_release>());
  m.def("get_cudnn_version", &transformer_engine::pytorch::get_cudnn_version, "Get cuDNN version",
        py::call_guard<py::gil_scoped_release>());
  m.def("copy_data_ptrs_to_device", &transformer_engine::pytorch::copy_data_ptrs_to_device,
        py::arg("tensors"), py::arg("device"), py::call_guard<py::gil_scoped_release>());
  m.def("splits_to_offsets", &transformer_engine::pytorch::splits_to_offsets,
        "Compute grouped tensor offsets from split sizes", py::arg("first_dims"),
        py::arg("logical_last_dim"), py::call_guard<py::gil_scoped_release>());
  m.def("splits_to_offsets_multi", &transformer_engine::pytorch::splits_to_offsets_multi,
        "Compute multiple scaled inclusive-scan offsets from a split-sizes vector",
        py::arg("split_sizes"), py::arg("device"), py::kw_only(), py::arg("strides"),
        py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false);
  m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams",
        py::call_guard<py::gil_scoped_release>());

  // Support THD format for Context Parallel
  m.def("thd_read_half_tensor", &transformer_engine::pytorch::thd_read_half_tensor,
        "Read the first half(half_idx=0) or the second half(half_idx=1) of each sequence in a THD "
        "tensor",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_second_half_lse_correction",
        &transformer_engine::pytorch::thd_second_half_lse_correction,
        "Correct the second half of the softmax_lse", py::call_guard<py::gil_scoped_release>());
  m.def("thd_read_second_half_lse", &transformer_engine::pytorch::thd_read_second_half_lse,
        "Read the second half of the softmax_lse", py::call_guard<py::gil_scoped_release>());
  m.def("thd_out_correction", &transformer_engine::pytorch::thd_out_correction,
        "Correct the THD format output of context parallelism in forward pass",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_grad_correction", &transformer_engine::pytorch::thd_grad_correction,
        "Correct the THD format gradients of context parallelism in backward pass",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_get_partitioned_indices", &transformer_engine::pytorch::thd_get_partitioned_indices,
        "Generate partitioned indices for inputs in THD format",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_sequence_order_to_cp_rank_order",
        &transformer_engine::pytorch::thd_sequence_order_to_cp_rank_order,
        "Reorder a THD tensor from sequence order to dual-chunk CP rank order",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_cp_rank_order_to_sequence_order",
        &transformer_engine::pytorch::thd_cp_rank_order_to_sequence_order,
        "Reorder a THD tensor from dual-chunk CP rank order to sequence order",
        py::call_guard<py::gil_scoped_release>());
  m.def("thd_copy_valid_tokens_from_per_split_to_rank_local",
        &transformer_engine::pytorch::thd_copy_valid_tokens_from_per_split_to_rank_local,
        "Copy valid THD token entries from a per-split tensor into a rank-local accumulator",
        py::call_guard<py::gil_scoped_release>());

  // nvshmem functions
  m.def("init_nvshmem_backend", &transformer_engine::pytorch::init_nvshmem_backend,
        "Initialize nvshmem backend with Pytorch distributed process groups",
        py::call_guard<py::gil_scoped_release>());
  m.def("create_nvshmem_tensor", &transformer_engine::pytorch::create_nvshmem_tensor,
        "Create a tensor in NVSHMEM shared memory", py::call_guard<py::gil_scoped_release>());
  m.def("nvshmem_send_on_current_stream",
        &transformer_engine::pytorch::nvshmem_send_on_current_stream,
        "Asynchronously send tensor data to a remote PE using NVSHMEM on the current CUDA stream",
        py::call_guard<py::gil_scoped_release>());
  m.def("nvshmem_wait_on_current_stream",
        &transformer_engine::pytorch::nvshmem_wait_on_current_stream,
        "Wait for a signal value to be updated by a remote PE using NVSHMEM on the current CUDA "
        "stream",
        py::call_guard<py::gil_scoped_release>());
  m.def("nvshmem_finalize", &transformer_engine::pytorch::nvshmem_finalize,
        "Clean up and finalize the NVSHMEM communication backend and free associated resources",
        py::call_guard<py::gil_scoped_release>());

  // multi-tensor functions
  m.def("multi_tensor_scale", &transformer_engine::pytorch::multi_tensor_scale_cuda,
        "Fused overflow check + scale for a list of contiguous tensors",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_scale_tensor", &transformer_engine::pytorch::multi_tensor_scale_tensor_cuda,
        "Fused overflow check + scale for a list of contiguous tensors with scale passed as tensor",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_l2norm", &transformer_engine::pytorch::multi_tensor_l2norm_cuda,
        "Computes L2 norm for a list of contiguous tensors",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_unscale_l2norm",
        &transformer_engine::pytorch::multi_tensor_unscale_l2norm_cuda,
        "Computes L2 norm for a list of contiguous tensors after unscaling (unscaling is only "
        "performed for L2 norm computation, and tensors are not updated)",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_adam", &transformer_engine::pytorch::multi_tensor_adam_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_adam_param_remainder",
        &transformer_engine::pytorch::multi_tensor_adam_param_remainder_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer"
        "where the master parameters only store the remainder bits",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_adam_fp8", &transformer_engine::pytorch::multi_tensor_adam_fp8_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_adam_capturable",
        &transformer_engine::pytorch::multi_tensor_adam_capturable_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer with CUDA graph "
        "support and LR scheduling",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_adam_capturable_master",
        &transformer_engine::pytorch::multi_tensor_adam_capturable_master_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer with CUDA graph "
        "support, LR scheduling and FP32 master weights",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_sgd", &transformer_engine::pytorch::multi_tensor_sgd_cuda,
        "Fused SGD optimizer for list of contiguous tensors",
        py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_compute_scale_and_scale_inv",
        &transformer_engine::pytorch::multi_tensor_compute_scale_and_scale_inv_cuda,
        "Fused compute scale and scale_inv from amax", py::call_guard<py::gil_scoped_release>());
  m.def("multi_tensor_compute_scale_inv_e8m0",
        &transformer_engine::pytorch::multi_tensor_compute_scale_inv_e8m0_cuda,
        "Fused compute E8M0 scale_inv from amax", py::call_guard<py::gil_scoped_release>());

  // Newton-Schulz (cuSolverMp)
  m.def("cusolvermp_ctx_create", &transformer_engine::pytorch::cusolvermp_ctx_create,
        "Create cuSolverMp context for Newton-Schulz", py::arg("nccl_comm_ptr"), py::arg("nranks"),
        py::arg("rank"), py::call_guard<py::gil_scoped_release>());
  m.def("cusolvermp_ctx_destroy", &transformer_engine::pytorch::cusolvermp_ctx_destroy,
        "Destroy cuSolverMp context", py::arg("ctx_ptr"), py::call_guard<py::gil_scoped_release>());
  m.def("newton_schulz", &transformer_engine::pytorch::newton_schulz,
        "Newton-Schulz matrix orthogonalization", py::arg("ctx_ptr"), py::arg("m"), py::arg("n"),
        py::arg("x"), py::arg("num_iterations"), py::arg("coefficients"),
        py::call_guard<py::gil_scoped_release>());

  // Comm+GEMM Overlap
  m.def("bulk_overlap_ag_with_external_gemm",
        &transformer_engine::pytorch::bulk_overlap_ag_with_external_gemm,
        "Bulk overlap All-Gather with a GEMM operation launched by another communicator",
        py::call_guard<py::gil_scoped_release>(), py::arg("allgather_communicator"),
        py::arg("send_stream"), py::arg("recv_stream"));

  // Experimental fused grouped MLP
  auto grouped_mlp_experimental = m.def_submodule(
      "grouped_mlp_experimental",
      "Experimental helpers for the fused grouped MLP (unstable, may change or disappear).");
  grouped_mlp_experimental.def("swizzle_scales_and_pack_ptrs_for_discrete_weights",
                               &transformer_engine::pytorch::grouped_mlp_experimental::
                                   swizzle_scales_and_pack_ptrs_for_discrete_weights,
                               py::arg("data_tensors"), py::arg("scale_tensors"),
                               py::arg("swizzle_type"), py::arg("device"),
                               py::call_guard<py::gil_scoped_release>());

  // Data structures
  py::class_<transformer_engine::pytorch::FP8TensorMeta>(m, "FP8TensorMeta")
      .def(py::init<>())
      .def_readwrite("scale", &transformer_engine::pytorch::FP8TensorMeta::scale)
      .def_readwrite("scale_inv", &transformer_engine::pytorch::FP8TensorMeta::scale_inv)
      .def_readwrite("amax_history", &transformer_engine::pytorch::FP8TensorMeta::amax_history);

  py::enum_<transformer_engine::pytorch::FP8FwdTensors>(m, "FP8FwdTensors")
      .value("GEMM1_INPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM1_INPUT)
      .value("GEMM1_WEIGHT", transformer_engine::pytorch::FP8FwdTensors::GEMM1_WEIGHT)
      .value("GEMM1_OUTPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM1_OUTPUT)
      .value("GEMM2_INPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM2_INPUT)
      .value("GEMM2_WEIGHT", transformer_engine::pytorch::FP8FwdTensors::GEMM2_WEIGHT)
      .value("GEMM2_OUTPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM2_OUTPUT)
      .value("GEMM3_INPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM3_INPUT)
      .value("GEMM3_WEIGHT", transformer_engine::pytorch::FP8FwdTensors::GEMM3_WEIGHT)
      .value("GEMM3_OUTPUT", transformer_engine::pytorch::FP8FwdTensors::GEMM3_OUTPUT);

  py::enum_<transformer_engine::pytorch::FP8BwdTensors>(m, "FP8BwdTensors")
      .value("GRAD_OUTPUT1", transformer_engine::pytorch::FP8BwdTensors::GRAD_OUTPUT1)
      .value("GRAD_INPUT1", transformer_engine::pytorch::FP8BwdTensors::GRAD_INPUT1)
      .value("GRAD_OUTPUT2", transformer_engine::pytorch::FP8BwdTensors::GRAD_OUTPUT2)
      .value("GRAD_INPUT2", transformer_engine::pytorch::FP8BwdTensors::GRAD_INPUT2)
      .value("GRAD_OUTPUT3", transformer_engine::pytorch::FP8BwdTensors::GRAD_OUTPUT3)
      .value("GRAD_INPUT3", transformer_engine::pytorch::FP8BwdTensors::GRAD_INPUT3);

  py::class_<CommOverlapHelper>(m, "CommOverlapHelper")
      .def(py::init<>(), py::call_guard<py::gil_scoped_release>())
      .def(py::init<c10d::ProcessGroup *, std::optional<c10d::ProcessGroup *>>(),
           py::call_guard<py::gil_scoped_release>(), py::arg("world_group"),
           py::arg("intra_node_group") = py::none());

  py::class_<CommOverlap, std::shared_ptr<CommOverlap>, transformer_engine::CommOverlapBase,
             transformer_engine::CommOverlapCore>(m, "CommOverlap")
      .def(py::init([](const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
                       CommOverlapHelper *helper, int tp_size, bool use_cublasmp,
                       transformer_engine::CommOverlapType comm_type, int num_splits,
                       int num_max_streams, int comm_cga_size, int gemm_priority, int comm_priority,
                       int num_comm_sm, bool set_sm_margin, bool atomic_gemm,
                       bool rs_overlap_first_gemm) {
             // Release the GIL only around the native construction (blocking collectives) to avoid
             // tripping pybind11's inc_ref/dec_ref GIL assertions.
             py::gil_scoped_release nogil;
             if (use_cublasmp) {
               return std::make_shared<CommOverlap>(helper, helper->mylocal, tp_size, comm_type,
                                                    buffer_shape, buffer_dtype, num_comm_sm,
                                                    atomic_gemm);
             }
             return std::make_shared<CommOverlap>(
                 buffer_shape, buffer_dtype, helper, tp_size, num_splits, num_max_streams,
                 comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin,
                 atomic_gemm, rs_overlap_first_gemm);
           }),
           py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"),
           py::arg("use_cublasmp") = false,
           py::arg("comm_type") = transformer_engine::CommOverlapType::RS,
           py::arg("num_splits") = 4, py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS,
           py::arg("comm_cga_size") = 2, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0,
           py::arg("num_comm_sm") = 16, py::arg("set_sm_margin") = true,
           py::arg("atomic_gemm") = false, py::arg("rs_overlap_first_gemm") = false)
      .def("copy_into_buffer",
           static_cast<void (CommOverlap::*)(const at::Tensor &, bool)>(
               &CommOverlap::copy_into_buffer),
           py::arg("input"), py::arg("local_chunk") = false)
      .def("get_buffer", &CommOverlap::get_buffer, py::arg("local_chunk") = false,
           py::arg("shape") = std::nullopt)
      .def("get_communication_stream", &CommOverlap::get_communication_stream);

  py::class_<CommOverlapP2P, std::shared_ptr<CommOverlapP2P>,
             transformer_engine::CommOverlapP2PBase, transformer_engine::CommOverlapCore>(
      m, "CommOverlapP2P")
      .def(py::init([](const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
                       CommOverlapHelper *helper, int tp_size,
                       transformer_engine::CommOverlapType comm_type, int num_max_streams,
                       int comm_cga_size, int gemm_priority, int comm_priority, int num_comm_sm,
                       bool set_sm_margin, bool atomic_gemm, bool use_ce, bool aggregate,
                       bool use_cublasmp) {
             // Release the GIL only around the native construction (blocking collectives) to avoid
             // tripping pybind11's inc_ref/dec_ref GIL assertions.
             py::gil_scoped_release nogil;
             if (use_cublasmp) {
               return std::make_shared<CommOverlapP2P>(helper, helper->mylocal, tp_size, comm_type,
                                                       buffer_shape, buffer_dtype, num_comm_sm,
                                                       atomic_gemm);
             }
             return std::make_shared<CommOverlapP2P>(buffer_shape, buffer_dtype, helper, tp_size,
                                                     comm_type, num_max_streams, comm_cga_size,
                                                     gemm_priority, comm_priority, num_comm_sm,
                                                     set_sm_margin, atomic_gemm, use_ce, aggregate);
           }),
           py::arg("buffer_shape"), py::arg("buffer_dtype"), py::arg("helper"), py::arg("tp_size"),
           py::arg("comm_type"), py::arg("num_max_streams") = NVTE_COMM_OVERLAP_MAX_STREAMS,
           py::arg("comm_cga_size") = 1, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0,
           py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false,
           py::arg("atomic_gemm") = false, py::arg("use_ce") = true, py::arg("aggregate") = false,
           py::arg("use_cublasmp") = false)
      .def("copy_into_buffer",
           static_cast<void (CommOverlapP2P::*)(const at::Tensor &, bool)>(
               &CommOverlapP2P::copy_into_buffer),
           py::arg("input"), py::arg("local_chunk") = false)
      .def("get_buffer", &CommOverlapP2P::get_buffer, py::arg("local_chunk") = false,
           py::arg("shape") = std::nullopt)
      .def("get_communication_stream", &CommOverlapP2P::get_communication_stream);
}  // NOLINT(readability/fn_size)
