/*
 * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "fp8_blockscale_gemm.h"
#ifdef ENABLE_FP8_BLOCK_SCALE
#include "fp8_blockscale_gemm_kernel.cuh"
#endif
#include "tensorrt_llm/common/logger.h"

namespace tensorrt_llm::kernels::fp8_blockscale_gemm {

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::gemm(
    void* mat_d, void const* mat_a, void const* mat_b, int shape_m, int shape_n, int shape_k,
    cudaStream_t stream, float const* scales_a, float const* scales_b) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  constexpr bool internal_quantize_a = !std::is_same_v<ElementA, __nv_fp8_e4m3>;
  constexpr bool internal_quantize_b = !std::is_same_v<ElementB, __nv_fp8_e4m3>;
  __nv_fp8_e4m3* fp8_mat_a;
  __nv_fp8_e4m3* fp8_mat_b;
  float* per_token_per_128c_scales;
  float* per_block_scales;

  auto* ws_ptr = workspace_;
  if constexpr (internal_quantize_a || internal_quantize_b) {
    TLLM_CHECK(ws_ptr != nullptr);
  }

  if constexpr (internal_quantize_a) {
    fp8_mat_a = reinterpret_cast<__nv_fp8_e4m3*>(ws_ptr);
    ws_ptr += max_shape_m_4_align_ * shape_k * sizeof(__nv_fp8_e4m3);
    per_token_per_128c_scales = reinterpret_cast<float*>(ws_ptr);
    ws_ptr += max_shape_m_4_align_ * div_up(shape_k, 128) * sizeof(float);
  }

  if constexpr (internal_quantize_b) {
    fp8_mat_b = reinterpret_cast<__nv_fp8_e4m3*>(ws_ptr);
    ws_ptr += shape_n * shape_k * sizeof(__nv_fp8_e4m3);
    per_block_scales = reinterpret_cast<float*>(ws_ptr);
    ws_ptr += div_up(shape_n, 128) * div_up(shape_k, 128) * sizeof(float);
  }

#ifdef COMPILE_HOPPER_TMA_GEMMS
  if constexpr (internal_quantize_a && internal_quantize_b) {
    fp8_gemm_run(reinterpret_cast<__nv_bfloat16 const*>(mat_a), fp8_mat_a, shape_k,
                 per_token_per_128c_scales, reinterpret_cast<__nv_bfloat16 const*>(mat_b),
                 fp8_mat_b, shape_k, per_block_scales, reinterpret_cast<__nv_bfloat16*>(mat_d),
                 shape_n, shape_m, shape_n, shape_k, stream, internal_quantize_a,
                 internal_quantize_b);
  }

  if constexpr (internal_quantize_a && !internal_quantize_b) {
    fp8_gemm_run(reinterpret_cast<__nv_bfloat16 const*>(mat_a), fp8_mat_a, shape_k,
                 per_token_per_128c_scales, nullptr,
                 reinterpret_cast<__nv_fp8_e4m3*>(const_cast<void*>(mat_b)), shape_k,
                 const_cast<float*>(scales_b), reinterpret_cast<__nv_bfloat16*>(mat_d), shape_n,
                 shape_m, shape_n, shape_k, stream, internal_quantize_a, internal_quantize_b);
  }
#else   // COMPILE_HOPPER_TMA_GEMMS
  TLLM_THROW("fp8 blockscale gemm only supported on Hopper.");
#endif  // COMPILE_HOPPER_TMA_GEMMS
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::gemm(
    __nv_fp8_e4m3 const* mat_a, int ld_a, __nv_fp8_e4m3 const* mat_b, int ld_b,
    __nv_bfloat16* mat_d, int ld_d, int shape_m, int shape_n, int shape_k, float const* scales_a,
    float const* scales_b, cudaStream_t stream) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  fp8_gemm_run(const_cast<__nv_fp8_e4m3*>(mat_a), ld_a, const_cast<__nv_fp8_e4m3*>(mat_b), ld_b,
               mat_d, ld_d, shape_m, shape_n, shape_k, const_cast<float*>(scales_a),
               const_cast<float*>(scales_b), stream);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::moeGemm(
    void* mat_d, void const* mat_a, void const* mat_b, int64_t const* problem_m_offsets,
    size_t num_problems, size_t shape_n, size_t shape_k, cudaStream_t stream, float const* scales_a,
    float const* scales_b) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  constexpr bool internal_quantize_a = !std::is_same_v<ElementA, __nv_fp8_e4m3>;
  constexpr bool internal_quantize_b = !std::is_same_v<ElementB, __nv_fp8_e4m3>;

  __nv_fp8_e4m3* fp8_mat_a;
  float* per_token_per_128c_scales;
  __nv_fp8_e4m3* fp8_mat_b;
  float* per_block_scales;

  auto* ws_ptr = workspace_;
  if constexpr (internal_quantize_a || internal_quantize_b) {
    TLLM_CHECK(ws_ptr != nullptr);
  }

  if constexpr (internal_quantize_a) {
    fp8_mat_a = reinterpret_cast<__nv_fp8_e4m3*>(ws_ptr);
    ws_ptr += max_shape_m_4_align_ * shape_k * sizeof(__nv_fp8_e4m3);
    per_token_per_128c_scales = reinterpret_cast<float*>(ws_ptr);
    ws_ptr += max_shape_m_32_align_padded_ * div_up(shape_k, 128) * sizeof(float);
  } else {
    fp8_mat_a = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<void*>(mat_a));
    per_token_per_128c_scales = const_cast<float*>(scales_a);
  }

  if constexpr (internal_quantize_b) {
    fp8_mat_b = reinterpret_cast<__nv_fp8_e4m3*>(ws_ptr);
    ws_ptr += num_problems * shape_n * shape_k * sizeof(__nv_fp8_e4m3);
    per_block_scales = reinterpret_cast<float*>(ws_ptr);
  } else {
    for (int i = 0; i < num_problems; i++) {
      fp8_mat_b = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<void*>(mat_b));
      per_block_scales = const_cast<float*>(scales_b);
    }
  }

#ifdef COMPILE_HOPPER_TMA_GEMMS
  if constexpr (std::is_same_v<ElementA, __nv_bfloat16> &&
                std::is_same_v<ElementB, __nv_bfloat16>) {
    fp8_grouped_gemm_run(reinterpret_cast<__nv_bfloat16 const*>(mat_a), fp8_mat_a,
                         per_token_per_128c_scales, reinterpret_cast<__nv_bfloat16 const*>(mat_b),
                         fp8_mat_b, per_block_scales, reinterpret_cast<__nv_bfloat16*>(mat_d),
                         problem_m_offsets, num_problems, expected_m_, max_shape_m_4_align_,
                         max_shape_m_32_align_padded_, shape_n, shape_k, stream,
                         internal_quantize_a, internal_quantize_b);
  } else if constexpr (std::is_same_v<ElementA, __nv_bfloat16> &&
                       std::is_same_v<ElementB, __nv_fp8_e4m3>) {
    fp8_grouped_gemm_run(reinterpret_cast<__nv_bfloat16 const*>(mat_a), fp8_mat_a,
                         per_token_per_128c_scales, nullptr, fp8_mat_b, per_block_scales,
                         reinterpret_cast<__nv_bfloat16*>(mat_d), problem_m_offsets, num_problems,
                         expected_m_, max_shape_m_4_align_, max_shape_m_32_align_padded_, shape_n,
                         shape_k, stream, internal_quantize_a, internal_quantize_b);
  } else if constexpr (std::is_same_v<ElementA, __nv_fp8_e4m3> &&
                       std::is_same_v<ElementB, __nv_fp8_e4m3>) {
    fp8_grouped_gemm_run(nullptr, fp8_mat_a, per_token_per_128c_scales,
                         reinterpret_cast<__nv_bfloat16 const*>(mat_b), fp8_mat_b, per_block_scales,
                         reinterpret_cast<__nv_bfloat16*>(mat_d), problem_m_offsets, num_problems,
                         expected_m_, max_shape_m_4_align_, max_shape_m_32_align_padded_, shape_n,
                         shape_k, stream, internal_quantize_a, internal_quantize_b);
  } else {
    TLLM_THROW("fp8 blockscale gemm only support __nv_fp8_e4m3 or bfloat16 as dataType.");
  }
#else   // COMPILE_HOPPER_TMA_GEMMS
  TLLM_THROW("fp8 blockscale gemm only support Hopper.");
#endif  // COMPILE_HOPPER_TMA_GEMMS
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::strideBatchGemm(
    __nv_bfloat16* mat_d, int ld_d, int stride_d, __nv_fp8_e4m3* mat_a, int ld_a, int stride_a,
    __nv_fp8_e4m3* mat_b, int ld_b, int stride_b, int num_problems, int shape_m, int shape_n,
    int shape_k, cudaStream_t stream, float* scales_a, int stride_scales_a, float* scales_b) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  fp8_stride_batch_gemm_run(nullptr, mat_a, scales_a, ld_a, stride_a, stride_scales_a, nullptr,
                            mat_b, scales_b, ld_b, stride_b, mat_d, ld_d, stride_d, num_problems,
                            shape_m, shape_n, shape_k, stream, false, false);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::fp8CS1x128(
    __nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y,
    cudaStream_t stream) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  fp8_1x128_cs(mat_quant, scales, mat, shape_x, shape_y, stream);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::fp8CS1x128Reshape(
    __nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_h,
    int shape_y, int stride_x, cudaStream_t stream) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  fp8_1x128_cs_reshape(mat_quant, scales, mat, shape_x, shape_h, shape_y, stride_x, stream);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
void CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::fp8CS128x128(
    __nv_fp8_e4m3* mat_quant, float* scales, __nv_bfloat16 const* mat, int shape_x, int shape_y,
    cudaStream_t stream) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  fp8_128x128_cs(mat_quant, scales, mat, shape_x, shape_y, stream);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getWorkspaceSizeBase(
    size_t max_shape_m, size_t shape_n, size_t shape_k, size_t num_problems) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  max_shape_m_4_align_ = std::max(max_shape_m_4_align_, int64_t(div_up(max_shape_m, 4) * 4));
  if (expected_m_ == 0) {
    expected_m_ = div_up(max_shape_m_4_align_, num_problems);
  }
  max_shape_m_32_align_padded_ = deep_gemm::compute_padded_offset(max_shape_m, num_problems);

  constexpr bool internal_quantize_a = !std::is_same_v<ElementA, __nv_fp8_e4m3>;
  constexpr bool internal_quantize_b = !std::is_same_v<ElementB, __nv_fp8_e4m3>;
  size_t total_workspace_size = 0;
  if constexpr (internal_quantize_a) {
    // fp8_mat_a
    total_workspace_size += max_shape_m_4_align_ * shape_k * sizeof(__nv_fp8_e4m3);
    // scales_a
    total_workspace_size += max_shape_m_32_align_padded_ * div_up(shape_k, 128) * sizeof(float);
  }

  if constexpr (internal_quantize_b) {
    // fp8_mat_b
    total_workspace_size += num_problems * shape_n * shape_k * sizeof(__nv_fp8_e4m3);
    // scales_b
    total_workspace_size +=
        num_problems * div_up(shape_k, 128) * div_up(shape_n, 128) * sizeof(float);
  }

  return total_workspace_size;
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getWorkspaceSize(
    size_t shape_m, size_t shape_n, size_t shape_k, size_t top_k, size_t num_problems) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  expected_m_ = shape_m;
  return getWorkspaceSizeBase(shape_m * top_k, shape_n, shape_k, num_problems);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getFP8DataSize(int shape_m,
                                                                                    int shape_n,
                                                                                    bool is_act) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  int shape_m_4_align = div_up(shape_m, 4) * 4;
  constexpr bool internal_quantize_a = !std::is_same_v<ElementA, __nv_fp8_e4m3>;
  constexpr bool internal_quantize_b = !std::is_same_v<ElementB, __nv_fp8_e4m3>;
  if (is_act && internal_quantize_a) {
    return div_up(shape_m_4_align * shape_n * sizeof(__nv_fp8_e4m3), 128) * 128;
  }

  if ((!is_act) && internal_quantize_b) {
    return div_up(shape_m * shape_n * sizeof(__nv_fp8_e4m3), 128) * 128;
  }
  return 0;
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getActScaleSize(int shape_m,
                                                                                     int shape_k) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  int shape_m_4_align = div_up(shape_m, 4) * 4;
  constexpr bool internal_quantize_a = !std::is_same_v<ElementA, __nv_fp8_e4m3>;
  size_t total_workspace_size = 0;
  if constexpr (internal_quantize_a) {
    // scales_a
    total_workspace_size +=
        div_up(shape_m_4_align * div_up(shape_k, 128) * sizeof(float), 128) * 128;
  }
  return total_workspace_size;
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getWeightScaleSize(
    int shape_n, int shape_k) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  constexpr bool internal_quantize_b = !std::is_same_v<ElementB, __nv_fp8_e4m3>;
  size_t total_workspace_size = 0;
  if constexpr (internal_quantize_b) {
    // scales_b
    total_workspace_size +=
        div_up(div_up(shape_k, 128) * div_up(shape_n, 128) * sizeof(float), 128) * 128;
  }

  return total_workspace_size;
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getActWorkspaceSize(
    int shape_m, int shape_k) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  return getFP8DataSize(shape_m, shape_k, true) + getActScaleSize(shape_m, shape_k);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template <typename ElementA, typename ElementB, typename ElementD>
size_t CutlassFp8BlockScaleGemmRunner<ElementA, ElementB, ElementD>::getWeightWorkspaceSize(
    int shape_n, int shape_k) {
#ifdef ENABLE_FP8_BLOCK_SCALE
  return getFP8DataSize(shape_n, shape_k, false) + getWeightScaleSize(shape_n, shape_k);
#else   // ENABLE_FP8_BLOCK_SCALE
  TLLM_THROW("fp8 blockscale gemm only supported on cuda version 12.8 or higher.");
  return 0;
#endif  // ENABLE_FP8_BLOCK_SCALE
}

template class CutlassFp8BlockScaleGemmRunner<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16>;
template class CutlassFp8BlockScaleGemmRunner<__nv_bfloat16, __nv_fp8_e4m3, __nv_bfloat16>;
template class CutlassFp8BlockScaleGemmRunner<__nv_fp8_e4m3, __nv_bfloat16, __nv_bfloat16>;
template class CutlassFp8BlockScaleGemmRunner<__nv_fp8_e4m3, __nv_fp8_e4m3, __nv_bfloat16>;

}  // namespace tensorrt_llm::kernels::fp8_blockscale_gemm
