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

#ifndef TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_
#define TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_

#include <cublasLt.h>
#include <cublas_v2.h>
#include <cuda_runtime_api.h>
#include <cudnn.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <transformer_engine/comm_gemm_overlap.h>
#include <transformer_engine/normalization.h>
#include <transformer_engine/transformer_engine.h>

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

#include "common/common.h"
#include "common/util/logging.h"
#include "extensions/ffi.h"
#include "extensions/misc.h"
#include "extensions/utils.h"
#include "transformer_engine/activation.h"
#include "transformer_engine/multi_stream.h"

namespace transformer_engine {
namespace jax {

struct ClampedSwigluConfig {
  float limit;
  float alpha;
  float glu_linear_offset;
};

struct ActivationConfig {
  ClampedSwigluConfig clamped_swiglu;
};

struct GemmConfig {
  JAXX_Scaling_Mode scaling_mode;
  JAXX_Collective_Op collective_op;
  int64_t lhs_axis_boundary;
  int64_t rhs_axis_boundary;
  bool lhs_transposed;
  bool rhs_transposed;
  bool use_split_accumulator;
};

struct GroupedGemmV2Config {
  bool lhs_is_trans;
  bool rhs_is_trans;
  JAXX_Scaling_Mode scaling_mode;
  int64_t lhs_axis_boundary;
  int64_t rhs_axis_boundary;
  int64_t lhs_left_size;
  int64_t lhs_right_size;
  int64_t rhs_left_size;
  int64_t rhs_right_size;
};

struct GroupedGemmConfig {
  bool lhs_is_trans;
  bool rhs_is_trans;
  JAXX_Scaling_Mode scaling_mode;
  bool has_bias;
  bool use_async_d2h_group_sizes;
  int64_t lhs_axis_boundary;
  int64_t rhs_axis_boundary;
  int64_t lhs_left_size;
  int64_t lhs_right_size;
  int64_t rhs_left_size;
  int64_t rhs_right_size;
};

inline bool use_fp8(DType type) { return type == DType::kFloat8E4M3 || type == DType::kFloat8E5M2; }

// Activation

XLA_FFI_DECLARE_HANDLER_SYMBOL(ActLuHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(ActLuInitializeHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(DActLuDBiasQuantizeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(DActLuDBiasQuantizeInitializeHandler);

pybind11::tuple GetDActDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size,
                                                   DType in_dtype, DType out_dtype,
                                                   JAXX_Scaling_Mode scaling_mode,
                                                   JAXX_Quantize_Layout quantize_layout);

// Normalization
XLA_FFI_DECLARE_HANDLER_SYMBOL(NormForwardInitializeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(NormForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(NormBackwardInitializeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(NormBackwardHandler);

pybind11::tuple GetNormForwardWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype,
                                             DType w_dtype, DType out_dtype,
                                             NVTE_Norm_Type norm_type,
                                             JAXX_Scaling_Mode scaling_mode,
                                             bool zero_centered_gamma, float epsilon, int sm_margin,
                                             bool is_training);

pybind11::tuple GetNormBackwardWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype,
                                              DType w_dtype, NVTE_Norm_Type norm_type,
                                              bool zero_centered_gamma, int sm_margin);

// Quantization
XLA_FFI_DECLARE_HANDLER_SYMBOL(DBiasQuantizeHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedQuantizeHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedQuantizeV2Handler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(DequantizeHandler);

pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size,
                                               DType in_dtype, DType out_dtype, DType scale_dtype,
                                               JAXX_Scaling_Mode scaling_mode,
                                               JAXX_Quantize_Layout quantize_layout);

// Softmax
XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledSoftmaxForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledSoftmaxBackwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledMaskedSoftmaxForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledMaskedSoftmaxBackwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledUpperTriangMaskedSoftmaxBackwardHandler);

// Attention
XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModForwardHandler);

XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnScoreModBackwardHandler);

NVTE_Fused_Attn_Backend GetFusedAttnBackend(
    bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout,
    NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type,
    float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen,
    size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left,
    int64_t window_size_right, bool deterministic);

pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
    size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
    size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim,
    size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
    NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout,
    DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left,
    int64_t window_size_right, bool bottom_right_diagonal);

pybind11::tuple GetFusedAttnBackwardWorkspaceSizes(
    size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
    size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim,
    size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
    NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout,
    DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq,
    int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal);

// GEMM
XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmV2Handler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(CollectiveGemmInitHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmInitV2Handler);

// Grouped GEMM
XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmD2HGroupSizesHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmV2Handler);

// Amax
XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationInitializeHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationHandler);

// Inspect
XLA_FFI_DECLARE_HANDLER_SYMBOL(InspectHandler);

// Cudnn helpers
XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnHandleInitHandler);

// CuBLAS helpers
XLA_FFI_DECLARE_HANDLER_SYMBOL(CublasHandleInitHandler);

// Router
XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossForwardHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler);

// Bootstrap EP (eager NCCL comm init); anchor released by ReleaseEpResources.
// max_token_dtype is the NVTEDType enum value (int) for the widest token dtype
// the group will dispatch.
void SetEpBootstrapParams(pybind11::bytes unique_id_bytes, int ep_size, int rank_within_group,
                          int num_experts, int max_tokens_per_rank, int max_recv_tokens_per_rank,
                          int hidden_dim, int max_num_sms, int max_token_dtype,
                          bool drop_on_overflow);
void ReleaseEpResources();
// Return the handle_mem byte size for a layer config.
size_t EpHandleMemSize(int top_k, size_t dispatch_output_per_expert_alignment);

// EpInstanceState type_id / type_info capsules for jax.ffi.register_ffi_type.
pybind11::capsule GetEpInstanceStateTypeIdCapsule();
pybind11::capsule GetEpInstanceStateTypeInfoCapsule();

// EP FFI handlers
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpInstantiateHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpPrepareHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpDispatchHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpDispatchBwdHandler);
XLA_FFI_DECLARE_HANDLER_SYMBOL(EpCombineBwdHandler);

// TopK
XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler);
pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k);

}  // namespace jax
}  // namespace transformer_engine

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(transformer_engine::jax::ClampedSwigluConfig,
                                      ::xla::ffi::StructMember<float>("limit"),
                                      ::xla::ffi::StructMember<float>("alpha"),
                                      ::xla::ffi::StructMember<float>("glu_linear_offset"));

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(
    transformer_engine::jax::ActivationConfig,
    ::xla::ffi::StructMember<transformer_engine::jax::ClampedSwigluConfig>("clamped_swiglu"));

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(
    transformer_engine::jax::GemmConfig,
    ::xla::ffi::StructMember<transformer_engine::jax::JAXX_Scaling_Mode>("scaling_mode"),
    ::xla::ffi::StructMember<transformer_engine::jax::JAXX_Collective_Op>("collective_op"),
    ::xla::ffi::StructMember<int64_t>("lhs_axis_boundary"),
    ::xla::ffi::StructMember<int64_t>("rhs_axis_boundary"),
    ::xla::ffi::StructMember<bool>("lhs_transposed"),
    ::xla::ffi::StructMember<bool>("rhs_transposed"),
    ::xla::ffi::StructMember<bool>("use_split_accumulator"));

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(
    transformer_engine::jax::GroupedGemmV2Config, ::xla::ffi::StructMember<bool>("lhs_is_trans"),
    ::xla::ffi::StructMember<bool>("rhs_is_trans"),
    ::xla::ffi::StructMember<transformer_engine::jax::JAXX_Scaling_Mode>("scaling_mode"),
    ::xla::ffi::StructMember<int64_t>("lhs_axis_boundary"),
    ::xla::ffi::StructMember<int64_t>("rhs_axis_boundary"),
    ::xla::ffi::StructMember<int64_t>("lhs_left_size"),
    ::xla::ffi::StructMember<int64_t>("lhs_right_size"),
    ::xla::ffi::StructMember<int64_t>("rhs_left_size"),
    ::xla::ffi::StructMember<int64_t>("rhs_right_size"));

XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(
    transformer_engine::jax::GroupedGemmConfig, ::xla::ffi::StructMember<bool>("lhs_is_trans"),
    ::xla::ffi::StructMember<bool>("rhs_is_trans"),
    ::xla::ffi::StructMember<transformer_engine::jax::JAXX_Scaling_Mode>("scaling_mode"),
    ::xla::ffi::StructMember<bool>("has_bias"),
    ::xla::ffi::StructMember<bool>("use_async_d2h_group_sizes"),
    ::xla::ffi::StructMember<int64_t>("lhs_axis_boundary"),
    ::xla::ffi::StructMember<int64_t>("rhs_axis_boundary"),
    ::xla::ffi::StructMember<int64_t>("lhs_left_size"),
    ::xla::ffi::StructMember<int64_t>("lhs_right_size"),
    ::xla::ffi::StructMember<int64_t>("rhs_left_size"),
    ::xla::ffi::StructMember<int64_t>("rhs_right_size"));

// ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace
XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode);
XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Score_Function);
XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Routing_Map_Format);
XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Collective_Op);
XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Quantize_Layout);

#endif  // TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_
