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

#ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_
#define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_

#include <nccl.h>

#include <cstdint>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include "common.h"

class CommOverlapHelper;
class CommOverlap;
class CommOverlapP2P;

namespace transformer_engine::pytorch {

/***************************************************************************************************
 * Router fusion
 **************************************************************************************************/

std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_topk_with_score_function_fwd(
    at::Tensor logits, int topk, bool use_pre_softmax, std::optional<int> num_groups,
    std::optional<int> group_topk, std::optional<float> scaling_factor, std::string score_function,
    std::optional<at::Tensor> expert_bias,
    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP),
    std::optional<at::Tensor> topk_indices = std::nullopt);

void fused_topk_with_score_function_bwd(
    at::Tensor routing_map, at::Tensor intermediate_output, at::Tensor grad_probs,
    at::Tensor grad_logits, int topk, bool use_pre_softmax, std::optional<float> scaling_factor,
    std::string score_function, bool use_dense_indices = false,
    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP));

std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_score_for_moe_aux_loss_fwd(
    at::Tensor logits, int topk, std::string score_function,
    int routing_map_format = static_cast<int>(NVTE_ROUTING_MAP_FORMAT_BYTEMAP));

void fused_score_for_moe_aux_loss_bwd(at::Tensor intermediate_output, at::Tensor grad_scores,
                                      at::Tensor grad_logits, int topk, std::string score_function);

std::tuple<at::Tensor, at::Tensor> fused_moe_aux_loss_fwd(at::Tensor probs,
                                                          at::Tensor tokens_per_expert,
                                                          int total_num_tokens, int num_experts,
                                                          int num_rows, int num_cols, int topk,
                                                          float coeff);

std::tuple<at::Tensor, at::Tensor> fused_moe_aux_loss_fwd_graph_safe(
    at::Tensor probs, at::Tensor tokens_per_expert, at::Tensor total_num_tokens, int num_experts,
    int num_rows, int num_cols, int topk, float coeff);

at::Tensor fused_moe_aux_loss_bwd(at::Tensor Const_buf, at::Tensor tokens_per_expert, int num_rows,
                                  int num_cols, at::Tensor grad_aux_loss);

/***************************************************************************************************
 * Permutation
 **************************************************************************************************/

std::tuple<at::Tensor, at::Tensor, std::vector<at::Tensor>> moe_permute_fwd(
    at::Tensor input, const DType dtype, at::Tensor indices, int64_t num_out_tokens,
    std::vector<at::Tensor> workspace, int64_t max_expanded_token_num);

at::Tensor moe_permute_bwd(at::Tensor input, const DType dtype, at::Tensor row_id_map,
                           at::Tensor prob, int64_t num_tokens, int64_t topK);

at::Tensor moe_unpermute_fwd(at::Tensor input, const DType dtype, at::Tensor row_id_map,
                             at::Tensor prob, int64_t num_tokens, int64_t topK);

std::tuple<at::Tensor, at::Tensor> moe_unpermute_bwd(at::Tensor input_bwd, at::Tensor input_fwd,
                                                     const DType dtype, at::Tensor row_id_map,
                                                     at::Tensor prob);

/***************************************************************************************************
 * Attention
 **************************************************************************************************/

NVTE_Fused_Attn_Backend get_fused_attn_backend(
    bool is_training, const DType q_dtype, const DType kv_dtype, NVTE_QKV_Layout qkv_layout,
    NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type,
    float p_dropout, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q,
    size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left,
    int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic);

std::vector<py::object> fused_attn_fwd(
    size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout,
    bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format,
    NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type,
    NVTE_Softmax_Type softmax_type, const std::vector<int64_t> window_size,
    bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv,
    const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype,
    const std::optional<at::Tensor> cu_seqlens_q_padded,
    const std::optional<at::Tensor> cu_seqlens_kv_padded,
    const std::optional<at::Tensor> page_table_k, const std::optional<at::Tensor> page_table_v,
    py::handle s_quantizer, py::handle o_quantizer, const std::optional<at::Tensor> Bias,
    const std::optional<at::Tensor> SoftmaxOffset, const std::optional<at::Generator> rng_gen,
    size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph);

std::vector<py::object> fused_attn_bwd(
    size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero,
    NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format,
    NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format,
    NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type,
    NVTE_Softmax_Type softmax_type, const std::vector<int64_t> window_size,
    bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q,
    const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V,
    const py::handle O, const py::handle dO, const at::ScalarType fake_dtype,
    const std::vector<at::Tensor> Aux_CTX_Tensors,
    const std::optional<at::Tensor> cu_seqlens_q_padded,
    const std::optional<at::Tensor> cu_seqlens_kv_padded, py::handle s_quantizer,
    py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph);

at::Tensor fa_prepare_fwd(at::Tensor qkvi);
at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v);

std::vector<std::optional<at::Tensor>> multi_tensor_transpose_to_bhsd(
    std::vector<std::optional<at::Tensor>> inputs, const std::string &original_format,
    std::vector<std::optional<at::Tensor>> outputs = {});

std::vector<at::Tensor> multi_tensor_pad_last_dim(std::vector<at::Tensor> inputs,
                                                  int64_t alignment);

at::Tensor convert_thd_to_bshd(at::Tensor tensor, at::Tensor cu_seqlens, int b, int max_seq_len);
at::Tensor convert_bshd_to_thd(at::Tensor tensor, at::Tensor cu_seqlens, int t);
void copy_to_kv_cache(at::Tensor new_k, at::Tensor new_v, at::Tensor k_cache, at::Tensor v_cache,
                      at::Tensor page_table, at::Tensor cu_new_lens, at::Tensor cu_cached_lens,
                      NVTE_QKV_Format kv_format, int b, int max_ctx_len, int max_seq_len,
                      int max_pages_per_seq, bool is_non_paged);

/***************************************************************************************************
 * GEMM
 **************************************************************************************************/

using MaybeTensor = std::optional<at::Tensor>;

std::vector<py::object> gemm(py::handle A, bool transa, py::handle B, bool transb, py::object D,
                             py::handle quantizer, std::optional<DType> out_dtype, MaybeTensor bias,
                             DType bias_type, bool gelu, MaybeTensor gelu_in, bool grad,
                             at::Tensor workspace, size_t workspaceSize, bool accumulate,
                             bool use_split_accumulator, CommOverlapCore *comm_overlap = nullptr,
                             std::optional<CommOverlapType> comm_type = std::nullopt,
                             MaybeTensor extra_output = std::nullopt, bool bulk_overlap = false,
                             float alpha = 1.0f, std::optional<float> beta = std::nullopt);

void te_atomic_gemm(at::Tensor A, at::Tensor A_scale_inverse, DType A_type,
                    std::vector<int64_t> A_scaling_mode, bool transa, at::Tensor B,
                    at::Tensor B_scale_inverse, DType B_type, std::vector<int64_t> B_scaling_mode,
                    bool transb, at::Tensor D, at::Tensor D_scale, DType D_type, at::Tensor D_amax,
                    at::Tensor bias, DType bias_type, at::Tensor pre_gelu_out, bool grad,
                    at::Tensor workspace, size_t workspaceSize, bool accumulate,
                    bool use_split_accumulator, int math_sm_count, int m_split, int n_split,
                    bool gemm_producer, at::Tensor counter);

std::optional<std::vector<at::Tensor>> te_general_grouped_gemm(
    std::vector<py::handle> A, bool transa, std::vector<py::handle> B, bool transb,
    std::optional<std::vector<at::Tensor>> D, DType D_type, std::vector<int64_t> m_splits,
    std::vector<at::Tensor> bias, DType bias_type, bool single_output,
    std::vector<at::Tensor> pre_gelu_out, bool grad, std::vector<at::Tensor> workspace,
    size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count);

py::object te_general_grouped_gemm_for_grouped_tensor(
    py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias,
    std::optional<at::Tensor> bias_scale, at::Tensor alpha, at::Tensor beta,
    at::Tensor workspace_setup, at::Tensor workspace_cublas, bool use_split_accumulator,
    int math_sm_count);

py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py::handle B,
                                                   bool transb, py::handle D, py::object bias,
                                                   std::optional<at::Tensor> bias_scale,
                                                   at::Tensor alpha, at::Tensor beta,
                                                   at::Tensor workspace_setup,
                                                   at::Tensor workspace_cublas,
                                                   bool use_split_accumulator, int math_sm_count);

py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, py::handle B,
                                                    bool transb, py::handle D, py::object bias,
                                                    std::optional<at::Tensor> bias_scale,
                                                    at::Tensor alpha, at::Tensor beta,
                                                    at::Tensor workspace_setup,
                                                    at::Tensor workspace_cublas,
                                                    bool use_split_accumulator, int math_sm_count);

/***************************************************************************************************
 * Transpose
 **************************************************************************************************/

at::Tensor fp8_transpose(at::Tensor input, DType otype,
                         std::optional<at::Tensor> output = std::nullopt);

at::Tensor nvfp4_data_transpose(at::Tensor input, std::optional<at::Tensor> output = std::nullopt);

void nvfp4_2d_scale_transpose(at::Tensor input, at::Tensor output, int64_t M_tiles,
                              int64_t K_tiles);

void nvfp4_2d_multi_tensor_transpose(std::vector<at::Tensor> rowwise_data_list,
                                     std::vector<at::Tensor> columnwise_data_list,
                                     std::vector<at::Tensor> rowwise_scale_inv_list,
                                     std::vector<at::Tensor> columnwise_scale_inv_list,
                                     std::vector<int64_t> M_list, std::vector<int64_t> K_list);

void nvfp4_multi_tensor_compute_partial_amax(
    std::vector<at::Tensor> master_weight_list, std::vector<at::Tensor> partial_amax_list,
    std::vector<at::Tensor> global_amax_list, std::vector<int64_t> h_list,
    std::vector<int64_t> w_list, std::vector<int64_t> start_offset_list, int64_t block_len);

void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows,
                               int64_t tile_cols, int64_t rows_padded, int64_t block_len);

void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax);

void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale,
                       at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows,
                       int64_t tile_cols, int64_t rows_padded, int64_t block_len);

void nvfp4_multi_tensor_fused_scale(
    std::vector<at::Tensor> block_amax_list, std::vector<at::Tensor> global_amax_list,
    std::vector<at::Tensor> per_block_scale_list, std::vector<at::Tensor> target_scale_list,
    std::vector<at::Tensor> target_amax_list, std::vector<int64_t> tile_rows_list,
    std::vector<int64_t> tile_cols_list, std::vector<int64_t> rows_padded_list, int64_t block_len);

void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale);

at::Tensor swap_first_dims(at::Tensor tensor, std::optional<at::Tensor> out = std::nullopt);

/***************************************************************************************************
 * Activations
 **************************************************************************************************/

/* GLU (sigmoid gate) */
py::object glu(const at::Tensor &input, py::handle quantizer);

py::object dglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

/* GELU and variants*/
py::object gelu(const at::Tensor &input, py::handle quantizer);

py::object dgelu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object geglu(const at::Tensor &input, py::handle quantizer);

py::object dgeglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object qgelu(const at::Tensor &input, py::handle quantizer);

py::object dqgelu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object qgeglu(const at::Tensor &input, py::handle quantizer);

py::object dqgeglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

/* ReLU and variants*/
py::object relu(const at::Tensor &input, py::handle quantizer);

py::object drelu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object reglu(const at::Tensor &input, py::handle quantizer);

py::object dreglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object srelu(const at::Tensor &input, py::handle quantizer);

py::object dsrelu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object sreglu(const at::Tensor &input, py::handle quantizer);

py::object dsreglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

/* Silu and variants*/
py::object silu(const at::Tensor &input, py::handle quantizer);

py::object dsilu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object swiglu(const at::Tensor &input, py::handle quantizer);

py::object dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer);

py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float limit, float alpha,
                          float glu_linear_offset);

py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer,
                           float limit, float alpha, float glu_linear_offset);

/* Scaled activation */
py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
                         py::handle quantizer, int64_t glu_interleave_size);

py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
                                 py::handle quantizer, float limit, float alpha,
                                 float glu_linear_offset, int64_t glu_interleave_size);

py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales,
                        py::handle quantizer);

py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input,
                         const at::Tensor &act_scales, py::handle quantizer,
                         int64_t glu_interleave_size, bool compute_scale_grad);

py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input,
                                 const at::Tensor &act_scales, py::handle quantizer, float limit,
                                 float alpha, float glu_linear_offset, int64_t glu_interleave_size,
                                 bool compute_scale_grad);

py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input,
                        const at::Tensor &act_scales, py::handle quantizer,
                        bool compute_scale_grad);

/***************************************************************************************************
 * LayerNorm
 **************************************************************************************************/

std::vector<py::object> layernorm_bwd(const at::Tensor &dz, const at::Tensor &x,
                                      const at::Tensor &mu, const at::Tensor &rsigma,
                                      const at::Tensor &gamma, const int sm_margin,
                                      const bool zero_centered_gamma);

std::vector<py::object> layernorm_fwd(py::handle input, py::handle weight, MaybeTensor bias,
                                      float eps, py::object ln_out, py::handle quantizer,
                                      DType out_dtype, const int sm_margin,
                                      const bool zero_centered_gamma);

/***************************************************************************************************
 * RMSNorm
 **************************************************************************************************/

std::vector<py::object> rmsnorm_bwd(const at::Tensor &dz, const at::Tensor &x,
                                    const at::Tensor &rsigma, const at::Tensor &gamma,
                                    const int sm_margin, const bool zero_centered_gamma);

std::vector<py::object> rmsnorm_bwd_add(const at::Tensor &dz, const at::Tensor &x,
                                        const at::Tensor &add, const at::Tensor &rsigma,
                                        const at::Tensor &gamma, const int sm_margin,
                                        const bool zero_centered_gamma);

std::vector<py::object> rmsnorm_fwd(const py::handle &input, const py::handle &weight, float eps,
                                    py::object ln_out, py::handle quantizer, DType otype,
                                    const int sm_margin, const bool zero_centered_gamma);

/***************************************************************************************************
 * Memory allocation
 **************************************************************************************************/

// Allocates tensors all backed by a single contiguous buffer.
std::vector<at::Tensor> bulk_allocate(const std::vector<std::vector<size_t>> &shapes,
                                      const std::vector<at::ScalarType> &dtypes,
                                      std::optional<c10::Device> device = std::nullopt,
                                      std::optional<std::vector<size_t>> alignments = std::nullopt);

/***************************************************************************************************
 * Quantize
 **************************************************************************************************/

py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector<size_t> &shape,
                                         at::ScalarType dtype, at::Device device, bool pin_memory);

py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output,
                    std::optional<at::Tensor> noop_flag);

py::object nvfp4_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                    const at::Tensor &rowwise_amax,
                                    const at::Tensor &columnwise_amax);

py::object dequantize(const py::handle &input, DType otype);

py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                          std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                          std::optional<at::Tensor> tensor_offsets,
                          std::optional<at::Tensor> noop_flag);

py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                          const size_t num_tensors,
                                          std::optional<at::Tensor> first_dims,
                                          std::optional<at::Tensor> last_dims,
                                          const at::Tensor &rowwise_amax,
                                          const at::Tensor &columnwise_amax,
                                          std::optional<at::Tensor> tensor_offsets);

py::object group_dequantize(const py::handle &input, DType otype);

py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                const size_t num_tensors, std::optional<at::Tensor> first_dims,
                                std::optional<at::Tensor> last_dims,
                                std::optional<at::Tensor> tensor_offsets);

py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
                                    DType otype, std::optional<at::Tensor> tensor_offsets,
                                    bool return_dequantized);

std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                              std::vector<py::handle> quantizer_list);

std::vector<py::object> split_quantize(const at::Tensor &tensor,
                                       const std::vector<size_t> &split_sections,
                                       std::vector<py::handle> quantizer_list,
                                       bool disable_bulk_allocation = false);

/***************************************************************************************************
 * Bias gradient fusions
 **************************************************************************************************/

std::vector<py::object> bgrad_quantize(const at::Tensor &input, py::handle py_quantizer);

std::vector<py::object> dbias_dgelu(const at::Tensor &grad_output, const at::Tensor &act_input,
                                    py::handle quantizer);

std::vector<py::object> dbias_dsilu(const at::Tensor &grad_output, const at::Tensor &act_input,
                                    py::handle quantizer);

std::vector<py::object> dbias_drelu(const at::Tensor &grad_output, const at::Tensor &act_input,
                                    py::handle quantizer);

std::vector<py::object> dbias_dqgelu(const at::Tensor &grad_output, const at::Tensor &act_input,
                                     py::handle quantizer);

std::vector<py::object> dbias_dsrelu(const at::Tensor &grad_output, const at::Tensor &act_input,
                                     py::handle quantizer);

/***************************************************************************************************
 * Dropout
 **************************************************************************************************/

std::vector<py::object> dropout_fwd(const py::handle &input, const float dropout_probability,
                                    std::optional<at::Tensor> out = std::nullopt);

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
 **************************************************************************************************/

void compute_amax(const at::Tensor &tensor, at::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);

// Note that the start_offset is the logical offset along the tensor dimension.
// The offset in bytes is start_offset * sizeof(tensor.dtype)
void fp8_block_scaling_compute_partial_amax(const at::Tensor &tensor, at::Tensor amax, size_t h,
                                            size_t w, size_t start_offset, size_t block_len);

void fp8_block_scaling_partial_cast(const at::Tensor &inp, at::Tensor out, const at::Tensor &scale,
                                    size_t h, size_t w, size_t start_offset, size_t block_len,
                                    const DType out_dtype);

void nvfp4_2d_compute_partial_amax(const at::Tensor &tensor, at::Tensor amax, size_t h, size_t w,
                                   size_t start_offset, size_t block_len);

void nvfp4_2d_partial_cast(const at::Tensor &inp, py::handle out, const at::Tensor &scale,
                           const at::Tensor &global_scale, size_t h, size_t w, size_t start_offset,
                           size_t block_len);

void nvfp4_multi_tensor_2d_partial_cast(std::vector<at::Tensor> inp_list,
                                        std::vector<at::Tensor> out_list,
                                        std::vector<at::Tensor> scale_list,
                                        std::vector<at::Tensor> global_scale_list,
                                        std::vector<int64_t> h_list, std::vector<int64_t> w_list,
                                        std::vector<int64_t> start_offset_list, int64_t block_len);
void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise,
                                        at::Tensor amax_colwise, int rows, int cols,
                                        size_t start_offset);

void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwise,
                                at::Tensor output_colwise, const at::Tensor &scale_inv_rowwise,
                                const at::Tensor &scale_inv_colwise, int rows, int cols,
                                size_t start_offset);

/***************************************************************************************************
 * Rotary positional embedding
 **************************************************************************************************/

at::Tensor fused_rope_forward(const at::Tensor &input, const at::Tensor &freqs,
                              const std::optional<at::Tensor> start_positions,
                              const NVTE_QKV_Format qkv_format, const bool interleaved,
                              const std::optional<at::Tensor> cu_seqlens, const int cp_size,
                              const int cp_rank);

at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor &freqs,
                               const std::optional<at::Tensor> start_positions,
                               const NVTE_QKV_Format qkv_format, const bool interleaved,
                               const std::optional<at::Tensor> cu_seqlens, const int cp_size,
                               const int cp_rank);

std::tuple<at::Tensor, at::Tensor, at::Tensor> fused_qkv_rope_forward(
    const at::Tensor &qkv_input, const at::Tensor &q_freqs, const at::Tensor &k_freqs,
    const std::optional<at::Tensor> start_positions, const std::vector<int> &qkv_split_arg_list,
    const NVTE_QKV_Format qkv_format, const bool interleaved, const int cp_size, const int cp_rank);

at::Tensor fused_qkv_rope_backward(const at::Tensor &q_grad_out, const at::Tensor &k_grad_out,
                                   const at::Tensor &v_grad_out, const at::Tensor &q_freqs,
                                   const at::Tensor &k_freqs,
                                   const std::vector<int> &qkv_split_arg_list,
                                   const NVTE_QKV_Format qkv_format, const bool interleaved,
                                   const int cp_size, const int cp_rank);

/***************************************************************************************************
 * Miscellaneous
 **************************************************************************************************/

size_t get_cublasLt_version();

size_t get_cudnn_version();

at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_dim);
std::tuple<at::Tensor, std::vector<at::Tensor>> splits_to_offsets_multi(
    const at::Tensor &split_sizes, const c10::Device &device, const std::vector<int64_t> &strides,
    const std::vector<bool> &include_leading_zero, const std::vector<at::ScalarType> &dtypes,
    bool bulk_allocate_outputs);

at::Tensor copy_data_ptrs_to_device(const std::vector<at::Tensor> &tensors,
                                    const c10::Device &device);

/***************************************************************************************************
 * Experimental helpers for the fused grouped MLP
 *
 * These primarily exist to support cuDNN CuTe DSL grouped GEMM
 * kernels. Since those are unstable and under active development,
 * these helpers should also be considered unstable.
 **************************************************************************************************/

namespace grouped_mlp_experimental {

// Prepare discrete weight tensors for the cuDNN CuTe DSL grouped GEMM
// kernel by swizzling scales and copying data and scale pointers to
// device. All tensors must share a uniform shape and `swizzle_type`
// must be one of "mxfp8_rowwise", "mxfp8_columnwise", or "nvfp4".
// Returns {data_ptrs_device, scale_ptrs_device, swizzled_scales_buffer}.
std::tuple<at::Tensor, at::Tensor, at::Tensor> swizzle_scales_and_pack_ptrs_for_discrete_weights(
    const std::vector<at::Tensor> &data_tensors, const std::vector<at::Tensor> &scale_tensors,
    const std::string &swizzle_type, const c10::Device &device);

}  // namespace grouped_mlp_experimental

/***************************************************************************************************
 * Support THD format for Context Parallel
 **************************************************************************************************/

at::Tensor thd_read_half_tensor(const at::Tensor &tensor, const at::Tensor &cu_seqlens,
                                int half_idx);

void thd_second_half_lse_correction(at::Tensor lse, const at::Tensor &lse_per_step,
                                    const at::Tensor &cu_seqlens, bool lse_packed);

at::Tensor thd_read_second_half_lse(const at::Tensor &lse, const at::Tensor &cu_seqlens,
                                    bool lse_packed, int second_half_lse_seqlen);

void thd_out_correction(at::Tensor out, const at::Tensor &out_per_step, const at::Tensor &lse,
                        const at::Tensor &lse_per_step, const at::Tensor &cu_seqlens,
                        bool only_second_half, bool lse_packed);

void thd_grad_correction(at::Tensor grad, const at::Tensor &grad_per_step,
                         const at::Tensor &cu_seqlens, const std::string &first_half,
                         const std::string &second_half);

at::Tensor thd_get_partitioned_indices(const at::Tensor &cu_seqlens, int total_tokens,
                                       int world_size, int rank);

at::Tensor thd_sequence_order_to_cp_rank_order(const at::Tensor &inp, const at::Tensor &cu_seqlens,
                                               int cp_size, int total_tokens);

at::Tensor thd_cp_rank_order_to_sequence_order(const at::Tensor &inp, const at::Tensor &cu_seqlens,
                                               int cp_size, int total_tokens);

void thd_copy_valid_tokens_from_per_split_to_rank_local(at::Tensor out, const at::Tensor &inp,
                                                        const at::Tensor &cu_seqlens_padded,
                                                        const at::Tensor &cu_seqlens);

/***************************************************************************************************
 * multi_tensor_* kernels
 **************************************************************************************************/

void multi_tensor_scale_cuda(int chunk_size, at::Tensor noop_flag,
                             std::vector<std::vector<at::Tensor>> tensor_lists, float scale);

void multi_tensor_scale_tensor_cuda(int chunk_size, at::Tensor is_infinite,
                                    std::vector<std::vector<at::Tensor>> tensor_lists,
                                    at::Tensor scale);

std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
    int chunk_size, at::Tensor noop_flag, std::vector<std::vector<at::Tensor>> tensor_lists,
    at::optional<bool> per_tensor_python);

std::tuple<at::Tensor, at::Tensor> multi_tensor_unscale_l2norm_cuda(
    int chunk_size, at::Tensor noop_flag, std::vector<std::vector<at::Tensor>> tensor_lists,
    at::Tensor inv_scale, at::optional<bool> per_tensor_python);

void multi_tensor_adam_cuda(int chunk_size, at::Tensor noop_flag,
                            std::vector<std::vector<at::Tensor>> tensor_lists, const float lr,
                            const float beta1, const float beta2, const float epsilon,
                            const int step, const int mode, const int bias_correction,
                            const float weight_decay);

void multi_tensor_adam_param_remainder_cuda(int chunk_size, at::Tensor noop_flag,
                                            std::vector<std::vector<at::Tensor>> tensor_lists,
                                            const float lr, const float beta1, const float beta2,
                                            const float epsilon, const int step, const int mode,
                                            const int bias_correction, const float weight_decay);

void multi_tensor_adam_fp8_cuda(int chunk_size, at::Tensor noop_flag,
                                std::vector<std::vector<at::Tensor>> tensor_lists, const float lr,
                                const float beta1, const float beta2, const float epsilon,
                                const int step, const int mode, const int bias_correction,
                                const float weight_decay, DType fp8_dtype);

void multi_tensor_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag,
                                       std::vector<std::vector<at::Tensor>> tensor_lists,
                                       at::Tensor lr, const float beta1, const float beta2,
                                       const float epsilon, at::Tensor step, const int mode,
                                       const int bias_correction, const float weight_decay,
                                       at::Tensor inv_scale);

void multi_tensor_adam_capturable_master_cuda(int chunk_size, at::Tensor noop_flag,
                                              std::vector<std::vector<at::Tensor>> tensor_lists,
                                              at::Tensor lr, const float beta1, const float beta2,
                                              const float epsilon, at::Tensor step, const int mode,
                                              const int bias_correction, const float weight_decay,
                                              at::Tensor inv_scale);

void multi_tensor_sgd_cuda(int chunk_size, at::Tensor noop_flag,
                           std::vector<std::vector<at::Tensor>> tensor_lists, float wd,
                           float momentum, float dampening, float lr, bool nesterov, bool first_run,
                           bool wd_after_momentum, float scale);

void multi_tensor_compute_scale_and_scale_inv_cuda(
    int chunk_size, at::Tensor noop_flag, std::vector<std::vector<at::Tensor>> tensor_lists,
    float max_fp8, bool force_pow_2_scales, float epsilon);

void multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, const py::object &dummy,
                                              std::vector<std::vector<at::Tensor>> tensor_lists);

/***************************************************************************************************
 * padding
 **************************************************************************************************/

void fused_multi_row_padding(at::Tensor input, at::Tensor output,
                             std::vector<size_t> input_row_list,
                             std::vector<size_t> padded_input_row_list);

void fused_multi_row_unpadding(at::Tensor input, at::Tensor output,
                               std::vector<size_t> input_row_list,
                               std::vector<size_t> unpadded_input_row_list);

/***************************************************************************************************
 * Scale swizzling for GEMM
 **************************************************************************************************/

void inplace_swizzle_scale_for_gemm(py::handle &tensor);

void inplace_multi_tensor_swizzle_scales_for_gemm(std::vector<py::object> &tensors,
                                                  bool rowwise_usage, bool columnwise_usage);

void inplace_multi_tensor_swizzle_scales_for_gemm_unchecked(std::vector<py::object> &tensors,
                                                            bool rowwise_usage,
                                                            bool columnwise_usage);

void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise);

/***************************************************************************************************
 * Expert Parallelism
 **************************************************************************************************/

// Borrows torch's NCCL host comm (``comm_ptr`` from ``ProcessGroupNCCL._comm_ptr()``).
// ``group_name`` is the PG name used by the symm-mem window resolver.
// ``zero_copy`` is forwarded into ``NVTEEpGroupConfig.zero_copy``.
void ep_initialize(uintptr_t comm_ptr, const std::string &group_name, int64_t num_experts,
                   int64_t max_tokens_per_rank, int64_t max_recv_tokens_per_rank,
                   int64_t hidden_dim, int64_t max_num_sms, pybind11::object max_token_dtype,
                   bool zero_copy, int64_t num_topk, bool drop_on_overflow);

void ep_finalize();

// Returns the zero-copy mode captured at ep_initialize.
bool ep_get_zero_copy();

// Returns the handle_mem byte size for the given layer config.
int64_t ep_handle_mem_size(int64_t top_k, int64_t dispatch_output_per_expert_alignment);

void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_per_expert,
                int64_t top_k, int64_t dispatch_output_per_expert_alignment,
                at::Tensor total_recv_tokens);

void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
                 at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights);

void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result);

void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights,
                     at::Tensor grad_tokens, at::Tensor grad_topk_weights);

void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out);

// Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP.
void register_ep_bindings(pybind11::module_ &m);

/***************************************************************************************************
 * NVSHMEM APIs
 **************************************************************************************************/

void init_nvshmem_backend(c10d::ProcessGroup *process_group);

at::Tensor create_nvshmem_tensor(const std::vector<int64_t> &shape, c10::ScalarType dtype);

void nvshmem_send_on_current_stream(at::Tensor src, at::Tensor dst, int peer, at::Tensor signal);

void nvshmem_wait_on_current_stream(at::Tensor signal, const std::string &wait_kind);

void nvshmem_finalize();

/***************************************************************************************************
 * Comm+GEMM Overlap Wrappers
 **************************************************************************************************/

void bulk_overlap_ag_with_external_gemm(CommOverlap &allgather_communicator, at::Stream send_stream,
                                        at::Stream recv_stream);

/***************************************************************************************************
 * Newton-Schulz (cuSolverMp)
 **************************************************************************************************/

int64_t cusolvermp_ctx_create(int64_t nccl_comm_ptr, int nranks, int rank);

void cusolvermp_ctx_destroy(int64_t ctx_ptr);

void newton_schulz(int64_t ctx_ptr, int64_t m, int64_t n, at::Tensor x, int64_t num_iterations,
                   std::vector<float> coefficients);

}  // namespace transformer_engine::pytorch

/***************************************************************************************************
 * Comm+GEMM Overlap Wrappers
 **************************************************************************************************/

class CommOverlapHelper : torch::CustomClassHolder {
 public:
  // Shared ownership of an ncclComm_t. The deleter calls ncclCommDestroy when
  // the last reference (held by the helper and/or any CommOverlap consumers)
  // is released, so the communicator outlives whichever owner is destroyed
  // first.
  using NcclCommSharedPtr = std::shared_ptr<std::remove_pointer<ncclComm_t>::type>;

 private:
  bool initialized{false};
  bool backend_is_nccl{false};
  std::map<std::string, c10d::ProcessGroup *> torch_pgs;
  std::map<std::string, NcclCommSharedPtr> nccl_comms;

 public:
  int myrank = -1;
  int numranks = -1;
  int mylocal = -1;
  int numlocal = -1;
  int mynode = -1;
  int numnodes = -1;

  CommOverlapHelper();

  CommOverlapHelper(c10d::ProcessGroup *world_group,
                    std::optional<c10d::ProcessGroup *> intra_node_group);

  ~CommOverlapHelper();

  void ub_allgather(void *globaldata, size_t globalbytes, void *localdata, size_t localbytes,
                    ExtComm comm);

  void ub_barrier(ExtComm comm);

  NcclCommSharedPtr get_nccl_comm(std::string comm_name);
};

class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOverlapBase {
 private:
  // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this
  // instance, independent of the CommOverlapHelper that created it.
  CommOverlapHelper::NcclCommSharedPtr _nccl_comm;

 public:
  CommOverlap(const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
              CommOverlapHelper *helper, int tp_size, int num_splits = 4,
              int num_max_streams = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 2,
              int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 16,
              bool set_sm_margin = true, bool atomic_gemm = false,
              bool rs_overlap_first_gemm = false);

  // cuBLASMp variant. `comm_type`, `buffer_shape`, and `buffer_dtype` size
  // the construction-time warmup matmul that primes cuBLASMp's lazy NCCL
  // window registrations and workspace allocation so subsequent matmuls
  // (including those captured in CUDA graphs) avoid the unsafe lazy paths.
  CommOverlap(CommOverlapHelper *helper, int tp_rank, int tp_size,
              transformer_engine::CommOverlapType comm_type,
              const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
              int num_comm_sm = 16, bool atomic_gemm = false);

  ~CommOverlap() {}

  using transformer_engine::CommOverlapCore::copy_into_buffer;
  void copy_into_buffer(const at::Tensor &input, bool local_chunk = false);

  at::Tensor get_buffer(bool local_chunk = false,
                        std::optional<std::vector<int64_t>> shape = std::nullopt);

  std::pair<at::Stream, at::Stream> get_communication_stream();

};  // CommOverlap

class CommOverlapP2P : torch::CustomClassHolder, public transformer_engine::CommOverlapP2PBase {
 private:
  // Keeps the cuBLASMp NCCL communicator alive for the lifetime of this
  // instance, independent of the CommOverlapHelper that created it.
  CommOverlapHelper::NcclCommSharedPtr _nccl_comm;

 public:
  CommOverlapP2P(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 = NVTE_COMM_OVERLAP_MAX_STREAMS, int comm_cga_size = 1,
                 int gemm_priority = 0, int comm_priority = 0, int num_comm_sm = 1,
                 bool set_sm_margin = false, bool atomic_gemm = false, bool use_ce = true,
                 bool aggregate = false);

  // cuBLASMp variant. See CommOverlap for the `comm_type`/buffer args.
  CommOverlapP2P(CommOverlapHelper *helper, int tp_rank, int tp_size,
                 transformer_engine::CommOverlapType comm_type,
                 const std::vector<size_t> &buffer_shape, at::ScalarType buffer_dtype,
                 int num_comm_sm = 1, bool atomic_gemm = false);

  ~CommOverlapP2P() {}

  using transformer_engine::CommOverlapP2PBase::copy_into_buffer;
  void copy_into_buffer(const at::Tensor &input, bool local_chunk = false);

  at::Tensor get_buffer(bool local_chunk = false,
                        std::optional<std::vector<int64_t>> shape = std::nullopt);

  std::pair<at::Stream, at::Stream> get_communication_stream();

};  // CommOverlapP2P

#endif  // TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_
