/* Copyright 2025 The xLLM Authors. 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

    https://github.com/jd-opensource/xllm/blob/main/LICENSE

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 "ops_api.h"

#if defined(USE_MLU)
#include "mlu/mlu_ops_api.h"
#elif defined(USE_NPU)
#include "core/kernels/npu/tilelang/tilelang_ops_api.h"
#include "npu/npu_ops_api.h"
#include "triton_npu/torch_api/triton_ops_api.h"
#elif defined(USE_CUDA)
#include "cuda/attention_runner.h"
#include "cuda/cuda_ops_api.h"
#elif defined(USE_ILU)
#include "ilu/ilu_ops_api.h"
#elif defined(USE_MUSA)
#include "cuda/cuda_ops_api.h"
#include "musa/musa_ops_api.h"
#endif

#include <numeric>

#include "common/macros.h"
#include "layers/common/attention_metadata.h"

namespace xllm::kernel {

void apply_rotary(RotaryParams& params) {
#if defined(USE_MLU)
  mlu::apply_rotary(params.q,
                    params.k,
                    params.sin,
                    params.cos,
                    params.position_ids,
                    params.cu_query_lens,
                    params.interleaved,
                    params.discrete,
                    params.dynamic_ntk,
                    params.max_query_len);
#elif defined(USE_NPU)
  npu::apply_rotary(
      params.q, params.k, params.cos_sin, params.position_ids.value());
#elif defined(USE_CUDA) || defined(USE_MUSA)
  bool is_neox = !params.interleaved;
  torch::Tensor pos_ids;
  torch::Tensor cos_sin;

  if (params.position_ids.has_value()) {
    // positions is already int64 on CUDA/MUSA (pre-converted in
    // ForwardInput::to).
    pos_ids = params.position_ids.value().to(torch::kInt64);
  } else if (params.cu_query_lens.has_value()) {
    auto cu = params.cu_query_lens.value().to(torch::kInt64);
    CHECK(cu.numel() >= 2) << "apply_rotary (CUDA): cu_query_lens must have at "
                              "least 2 elements when "
                              "position_ids is not provided.";
    int64_t seq_len = cu[1].item<int64_t>() - cu[0].item<int64_t>();
    CHECK(seq_len > 0)
        << "apply_rotary (CUDA): invalid sequence length inferred from "
           "cu_query_lens when position_ids is not provided.";
    pos_ids = torch::arange(seq_len,
                            torch::TensorOptions()
                                .dtype(torch::kInt64)
                                .device(params.q.device()))
                  .contiguous();
  } else {
    // When neither position_ids nor cu_query_lens is provided,
    // infer sequence length from q tensor and create default position IDs.
    // This handles cases like LongCat-Image-Edit where rotary embedding
    // is applied uniformly across all sequence positions.
    int64_t seq_len = params.q.size(0);
    CHECK(seq_len > 0) << "apply_rotary (CUDA): cannot infer valid sequence "
                          "length from q tensor.";
    pos_ids = torch::arange(seq_len,
                            torch::TensorOptions()
                                .dtype(torch::kInt64)
                                .device(params.q.device()))
                  .contiguous();
  }

  if (params.precomputed_cos_sin.defined()) {
    cos_sin = params.precomputed_cos_sin;
  } else if (params.cos.defined() && params.sin.defined()) {
    const int64_t head_dim = params.cos.size(-1);
    const int64_t rot_half = head_dim / 2;
    auto cos_sliced = params.cos.contiguous().slice(-1, 0, rot_half);
    auto sin_sliced = params.sin.contiguous().slice(-1, 0, rot_half);
    cos_sin = torch::cat({cos_sliced, sin_sliced}, -1);
  } else if (params.cos_sin.defined()) {
    auto cos_sin_vec = params.cos_sin.chunk(4, -1);
    auto cos = cos_sin_vec[0];
    auto sin = cos_sin_vec[2];
    cos_sin = torch::cat({cos, sin}, -1);
  } else {
    LOG(FATAL) << "apply_rotary (CUDA): neither cos_sin nor cos/sin "
                  "provided; cannot infer cos_sin.";
  }

  cuda::rotary_embedding(pos_ids, params.q, params.k, cos_sin, is_neox);
#elif defined(USE_ILU)
  torch::Tensor ilu_cos_sin;
  if (params.precomputed_cos_sin.defined()) {
    ilu_cos_sin = params.precomputed_cos_sin;
  } else {
    auto cos_sin_vec = params.cos_sin.chunk(4, -1);
    ilu_cos_sin = torch::cat({cos_sin_vec[0], cos_sin_vec[2]}, -1);
  }
  // positions is already int64 on ILU (pre-converted in ForwardInput::to).
  torch::Tensor long_position_ids = params.position_ids.value().to(at::kLong);
  ilu::apply_rope_pos_ids_cos_sin_cache(
      params.q, params.k, ilu_cos_sin, long_position_ids, params.interleaved);
#else
  NOT_IMPLEMENTED();
#endif
}

void active(ActivationParams& params) {
#if defined(USE_MLU)
  mlu::active(params.input,
              params.output,
              params.bias,
              params.cusum_token_count,
              params.act_mode,
              params.is_gated,
              params.start_expert_id,
              params.expert_size);
#elif defined(USE_NPU)
  params.output = npu::active(params.input, params.act_mode);
#elif defined(USE_CUDA) || defined(USE_MUSA)
  cuda::act_and_mul(params.output, params.input, params.act_mode);
#elif defined(USE_ILU)
  ilu::act_and_mul(params.output, params.input, params.act_mode);
#else
  NOT_IMPLEMENTED();
#endif
}

void reshape_paged_cache(ReshapePagedCacheParams& params) {
#if defined(USE_MLU)
  mlu::reshape_paged_cache(params.key,
                           params.value,
                           params.k_cache,
                           params.v_cache,
                           params.slot_mapping,
                           params.direction);
#elif defined(USE_NPU)
  npu::reshape_paged_cache(params.key,
                           params.value,
                           params.k_cache,
                           params.v_cache,
                           params.slot_mapping);
#elif defined(USE_CUDA) || defined(USE_MUSA)
  cuda::reshape_paged_cache(params.slot_mapping,
                            params.key,
                            params.value.value_or(torch::Tensor()),
                            params.k_cache,
                            params.v_cache.value_or(torch::Tensor()));
#elif defined(USE_ILU)
  // auto v_cache = params.v_cache.value_or(torch::Tensor());
  ilu::reshape_paged_cache(params.key,
                           params.value,
                           params.k_cache,
                           params.v_cache,
                           params.slot_mapping);
#else
  NOT_IMPLEMENTED();
#endif
}

void reshape_from_cache(ReshapeFromCacheParams& params) {
#if defined(USE_MLU)
  mlu::reshape_from_cache(params.key,
                          params.value,
                          params.key_cache,
                          params.value_cache,
                          params.context_lengths,
                          params.max_context_len,
                          params.context_seq_offset,
                          params.block_tables,
                          params.cache_seq_offset);
#else
  NOT_IMPLEMENTED();
#endif
}

void quant_to_paged_cache(ReshapePagedCacheParams& params) {
#if defined(USE_MLU)
  CHECK(params.k_cache_scale.has_value())
      << "k_cache_scale is required for quant_to_paged_cache";
  mlu::quant_to_paged_cache(params.key,
                            params.value,
                            params.k_cache,
                            params.v_cache,
                            params.k_cache_scale.value(),
                            params.v_cache_scale,
                            params.slot_mapping);
#else
  NOT_IMPLEMENTED();
#endif
}

void dequant_from_paged_cache(ReshapeFromCacheParams& params) {
#if defined(USE_MLU)
  CHECK(params.key_cache_quant_scale.has_value())
      << "key_cache_quant_scale is required for dequant_from_paged_cache";
  mlu::dequant_from_paged_cache(params.key,
                                params.value,
                                params.key_cache,
                                params.value_cache,
                                params.key_cache_quant_scale.value(),
                                params.value_cache_quant_scale,
                                params.context_lengths,
                                params.max_context_len,
                                params.context_seq_offset,
                                params.block_tables.value(),
                                params.quant_mode,
                                params.quant_bit);
#else
  NOT_IMPLEMENTED();
#endif
}

void fused_layernorm(FusedLayerNormParams& params) {
#if defined(USE_MLU)
  mlu::fused_layernorm(params.input,
                       params.output,
                       params.residual,
                       params.weight,
                       params.beta,
                       params.bias,
                       params.quant_scale,
                       params.residual_out,
                       params.smooth_quant_scale,
                       params.normed_out,
                       params.mode,
                       params.eps,
                       params.store_output_before_norm,
                       params.store_output_after_norm,
                       params.dynamic_quant);
#elif defined(USE_MUSA)
  musa::fused_layernorm(params.input,
                        params.output,
                        params.residual,
                        params.weight,
                        params.beta,
                        params.bias,
                        params.quant_scale,
                        params.residual_out,
                        params.smooth_quant_scale,
                        params.normed_out,
                        params.mode,
                        params.eps,
                        params.store_output_before_norm,
                        params.store_output_after_norm,
                        params.dynamic_quant);
#elif defined(USE_NPU)
  if (params.residual.has_value()) {
    std::tie(params.output, std::ignore, params.residual_out) =
        npu::add_rms_norm(
            params.input, params.residual.value(), params.weight, params.eps);
  } else {
    params.output =
        npu::rms_norm(params.input, params.weight, params.eps, params.mode);
  }
#elif defined(USE_CUDA) || defined(USE_MUSA)
  if (params.residual.has_value()) {
    cuda::fused_add_rms_norm(
        params.input, params.residual.value(), params.weight, params.eps);
    params.output = params.input;
    params.residual_out = params.residual;
  } else {
    cuda::rms_norm(params.output, params.input, params.weight, params.eps);
  }
#elif defined(USE_ILU)
  if (params.residual.has_value()) {
    ilu::residual_layer_norm(params.input,
                             params.output,
                             params.residual,
                             params.weight,
                             params.bias,  // residual_bias
                             params.residual_out,
                             params.eps);
  } else {
    ilu::rms_norm(params.output, params.input, params.weight, params.eps);
  }
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor matmul(MatmulParams& params) {
#if defined(USE_MLU)
  return mlu::matmul(
      params.a, params.b, params.bias, params.c, params.alpha, params.beta);
#elif defined(USE_NPU)
  return npu::matmul(params.a, params.b, params.bias);
#elif defined(USE_CUDA) || defined(USE_MUSA)
  return cuda::matmul(params.a, params.b, params.bias);
#elif defined(USE_ILU)
  return ilu::matmul(params.a, params.b, params.bias);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor group_gemm(GroupGemmParams& params) {
#if defined(USE_MLU)
  return mlu::group_gemm(params.a,
                         params.b,
                         params.token_count,
                         params.output,
                         params.a_scale,
                         params.b_scale,
                         params.quant_flag,
                         params.max_dim,
                         params.trans_a,
                         params.trans_b,
                         params.a_quant_bit);
#elif defined(USE_NPU)
  std::vector<torch::Tensor> x_list;
  std::vector<torch::Tensor> weight_list;
  torch::TensorList x_ref;
  torch::TensorList weight_ref;
  if (params.x_list.has_value()) {
    x_ref = params.x_list.value();
  } else {
    x_list = {params.a};
    x_ref = x_list;
  }
  if (params.weight_list.has_value()) {
    weight_ref = params.weight_list.value();
  } else {
    weight_list = {params.b};
    weight_ref = weight_list;
  }
  std::optional<torch::Tensor> group_list = params.group_list;
  if (!group_list.has_value()) {
    group_list = params.token_count;
  }

  auto outputs =
      npu::apply_npu_grouped_matmul(x_ref,
                                    weight_ref,
                                    params.bias_list,
                                    params.scale_list,
                                    params.offset_list,
                                    params.antiquant_scale_list,
                                    params.antiquant_offset_list,
                                    params.per_token_scale_list,
                                    group_list,
                                    params.activation_input_list,
                                    params.activation_quant_scale_list,
                                    params.activation_quant_offset_list,
                                    params.split_item,
                                    params.group_type,
                                    params.group_list_type,
                                    params.act_type,
                                    params.tuning_config,
                                    params.output_dtype);
  return outputs.back();
#elif defined(USE_ILU)
  return ilu::group_gemm(params.a,
                         params.b,
                         params.token_count,
                         params.combine_idx,
                         params.output);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
    MoeFusedTopkParams& params) {
#if defined(USE_MLU)
  return mlu::moe_active_topk(params.input,
                              params.topk,
                              params.num_expert_group,
                              params.topk_group,
                              params.normalize,
                              params.mask,
                              params.normed_by,
                              params.scoring_func,
                              params.route_scale,
                              params.e_score_correction_bias);
#elif defined(USE_NPU)
  CHECK_EQ(params.scoring_func, "softmax")
      << "Only softmax is supported for NPU";
  auto [topk_weights, topk_ids, row_ids] = npu::apply_moe_gating_topk_softmax(
      params.input, params.finished, params.topk);
  (void)row_ids;
  return std::make_tuple(topk_weights, topk_ids);
#elif defined(USE_ILU)
  return ilu::moe_active_topk(params.input,
                              params.topk,
                              params.num_expert_group,
                              params.topk_group,
                              params.normalize,
                              params.mask,
                              params.normed_by,
                              params.scoring_func,
                              params.route_scale,
                              params.e_score_correction_bias);
#elif defined(USE_CUDA) || defined(USE_MUSA)
  return cuda::moe_fused_topk(params.input,
                              params.topk,
                              params.normalize,
                              params.e_score_correction_bias,
                              params.scoring_func);
#else
  NOT_IMPLEMENTED();
#endif
}

std::vector<torch::Tensor> moe_gen_idx(MoeGenIdxParams& params) {
#if defined(USE_MLU)
  return mlu::moe_gen_idx(params.expert_id, params.expert_num);
#elif defined(USE_ILU)
  return ilu::moe_gen_idx(params.expert_id, params.expert_num);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor moe_expand_input(MoeExpandInputParams& params) {
#if defined(USE_MLU)
  return mlu::moe_expand_input(params.input,
                               params.gather_index,
                               params.cusum_token_count,
                               params.start_expert_id,
                               params.expert_size);
#elif defined(USE_ILU)
  return ilu::moe_expand_input(
      params.input, params.gather_index, params.combine_idx, params.topk);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor moe_combine_result(MoeCombineResultParams& params) {
#if defined(USE_MLU)
  return mlu::moe_combine_result(params.input,
                                 params.reduce_weight,
                                 params.gather_ids,
                                 params.residual,
                                 params.cusum_token_count,
                                 params.start_expert_id,
                                 params.expert_size,
                                 params.bias);
#elif defined(USE_NPU)
  std::optional<torch::Tensor> probes =
      params.probes.has_value()
          ? params.probes
          : std::optional<torch::Tensor>(params.reduce_weight);
  auto output = npu::apply_npu_moe_token_unpermute(params.input,
                                                   params.gather_ids,
                                                   probes,
                                                   params.padded_mode,
                                                   params.restore_shape);
  if (params.residual.has_value()) {
    output = output + params.residual.value();
  }
  return output;
#elif defined(USE_ILU)
  return ilu::moe_combine_result(params.input, params.reduce_weight);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor moe_all2all_gen_send_layout(
    MoeAll2AllGenSendLayoutParams& params) {
#if defined(USE_MLU)
  return mlu::moe_all2all_gen_send_layout(params.token_count, params.nrank);
#else
  NOT_IMPLEMENTED();
#endif
}

std::vector<torch::Tensor> moe_all2all_gen_gather_index(
    MoeAll2AllGenGatherIndexParams& params) {
#if defined(USE_MLU)
  return mlu::moe_all2all_gen_gather_index(
      params.token_num, params.pad_num, params.return_cusum_token_count);
#else
  NOT_IMPLEMENTED();
#endif
}

std::vector<torch::Tensor> moe_all2all_create(MoeAll2AllCreateParams& params) {
#if defined(USE_MLU)
  return mlu::moe_all2all_create(params.dispatch_token_byte,
                                 params.combine_token_byte,
                                 params.max_expert_num,
                                 params.max_token_num,
                                 params.rank,
                                 params.nrank,
                                 params.device);
#else
  NOT_IMPLEMENTED();
#endif
}

void moe_all2all_init(MoeAll2AllInitParams& params) {
#if defined(USE_MLU)
  mlu::moe_all2all_init(params.handle, params.all_exchange_info, params.device);
#else
  NOT_IMPLEMENTED();
#endif
}

void moe_all2all_dispatch(MoeAll2AllDispatchParams& params) {
#if defined(USE_MLU)
  mlu::moe_all2all_dispatch(params.handle,
                            params.token_byte,
                            params.token_num,
                            params.send_layout,
                            params.send_token_num,
                            params.recv_layout,
                            params.recv_token_num,
                            params.send_token,
                            params.recv_token);
#else
  NOT_IMPLEMENTED();
#endif
}

void moe_all2all_combine(MoeAll2AllCombineParams& params) {
#if defined(USE_MLU)
  mlu::moe_all2all_combine(params.handle,
                           params.token_byte,
                           params.token_num,
                           params.send_src_layout,
                           params.send_dst_layout,
                           params.send_token,
                           params.recv_token);
#else
  NOT_IMPLEMENTED();
#endif
}

void moe_all2all_destroy(MoeAll2AllDestroyParams& params) {
#if defined(USE_MLU)
  mlu::moe_all2all_destroy(params.handle, params.device);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor> scaled_quantize(
    ScaledQuantizeParams& params) {
#if defined(USE_MLU)
  return mlu::scaled_quantize(params.x,
                              params.smooth,
                              params.zero,
                              params.token_count,
                              params.gather_index,
                              params.gather_index_start_position,
                              params.output,
                              params.output_scale,
                              params.act_mode,
                              params.active_coef,
                              params.is_gated,
                              params.quant_type);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor scaled_matmul(ScaledMatmulParams& params) {
#if defined(USE_MLU)
  return mlu::scaled_matmul(params.a,
                            params.b,
                            params.a_scale,
                            params.b_scale,
                            params.output_dtype,
                            params.bias,
                            params.c,
                            params.act_mode,
                            params.quant_bit_size,
                            params.alpha,
                            params.beta,
                            params.use_hp_active,
                            params.a_quant_bit_size,
                            params.a_calib,
                            params.b_calib,
                            params.output);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor apply_top_k_top_p(TopKPParams& params) {
#if defined(USE_MLU)
  return mlu::apply_top_k_top_p(
      params.logits, params.temperatures, params.top_k, params.top_p);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor random_sample(RandomSampleParams& params) {
#if defined(USE_MLU)
  return mlu::random_sample(params.logits);
#elif defined(USE_CUDA)
  return cuda::random_sample(params.logits);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor rejection_sample(RejectionSampleParams& params) {
#if defined(USE_MLU)
  return mlu::rejection_sample(params.draft_token_ids,
                               params.num_draft_tokens,
                               params.cu_num_draft_tokens,
                               params.draft_probs,
                               params.target_probs,
                               params.bonus_token_ids,
                               params.uniform_rand,
                               params.uniform_probs,
                               params.max_spec_len);
#else
  NOT_IMPLEMENTED();
#endif
}

void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params) {
#if defined(USE_MLU)
  mlu::masked_indexer_select_paged_kv(params.query,
                                      params.k_cache,
                                      params.weights,
                                      params.kv_cache_block_table,
                                      params.cu_seq_q_lens,
                                      params.cu_seq_k_lens,
                                      params.k_context_lens,
                                      params.k_cache_block_table,
                                      params.is_prefill,
                                      params.index_topk,
                                      params.kv_cache_block_size,
                                      params.softmax_scale,
                                      params.q_scale,
                                      params.k_scale_cache,
                                      params.sparse_block_table,
                                      params.sparse_context_lens);
#else
  NOT_IMPLEMENTED();
#endif
}

void gather_split(GatherSplitParams& params) {
#if defined(USE_MLU)
  mlu::gather_split(params.input,
                    params.gather_index,
                    params.valid_token_num,
                    params.output_head,
                    params.output_tail);
#else
  NOT_IMPLEMENTED();
#endif
}

void fused_mla_q(FusedMlaQParams& params) {
#if defined(USE_MLU)
  mlu::fused_mla_q(params.q,
                   params.output,
                   params.output_scale,
                   params.output_norm,
                   params.gamma,
                   params.smooth_quant_scale,
                   params.weight_b,
                   params.weight_b_scale,
                   params.weight_c,
                   params.sin,
                   params.cos,
                   params.position_id,
                   params.quant_mode,
                   params.eps,
                   params.interleaved);
#else
  NOT_IMPLEMENTED();
#endif
}

void fused_mla_kv(FusedMlaKVParams& params) {
#if defined(USE_MLU)
  mlu::fused_mla_kv(params.input_kv,
                    params.sin,
                    params.cos,
                    params.position_id,
                    params.gamma,
                    params.kv_cache,
                    params.kv_cache_scale,
                    params.slot_mapping,
                    params.cache_bs_id,
                    params.cache_seq_offset,
                    params.quant_mode,
                    params.is_paged_cache,
                    params.eps,
                    params.interleaved);
#else
  NOT_IMPLEMENTED();
#endif
}

void fused_indexer_q(FusedIndexerQParams& params) {
#if defined(USE_MLU)
  mlu::fused_indexer_q(params.input_q,
                       params.output,
                       params.output_scale,
                       params.w_q,
                       params.w_q_scale,
                       params.hadamard_matrix,
                       params.sin,
                       params.cos,
                       params.position_id,
                       params.quant_mode,
                       params.interleaved,
                       params.rope_at_front);
#else
  NOT_IMPLEMENTED();
#endif
}

void fused_indexer_k(FusedIndexerKParams& params) {
#if defined(USE_MLU)
  mlu::fused_indexer_k(params.x,
                       params.wk,
                       params.wproj,
                       params.sin_table,
                       params.cos_table,
                       params.position_id,
                       params.slot_mapping,
                       params.head_weights,
                       params.k_cache,
                       params.k_cache_scale,
                       params.hadamard_matrix,
                       params.interleaved,
                       params.gamma,
                       params.beta,
                       params.eps);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor l2_norm(torch::Tensor& x, double eps) {
#if defined(USE_NPU)
  return npu::npu_l2norm_last_dim(x, eps);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
moe_init_routing_v2(MoeInitRoutingV2Params& params) {
#if defined(USE_NPU)
  return npu::apply_npu_moe_init_routing_v2(params.x,
                                            params.expert_idx,
                                            params.scale,
                                            params.offset,
                                            params.active_num,
                                            params.expert_capacity,
                                            params.expert_num,
                                            params.drop_pad_mode,
                                            params.expert_tokens_num_type,
                                            params.expert_tokens_num_flag,
                                            params.quant_mode,
                                            params.active_expert_range,
                                            params.row_idx_type);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
    Fp8ScaledQuantizeParams& params) {
#if defined(USE_CUDA)
  return cuda::fp8_scaled_quantize(params.input, params.output, params.scale);
#else
  NOT_IMPLEMENTED();
#endif
}

std::pair<torch::Tensor, torch::Tensor> fused_gdn_gating(
    FusedGdnGatingParams& params) {
#if defined(USE_NPU)
  return npu::tilelang::fused_gdn_gating(params.A_log,
                                         params.a,
                                         params.b,
                                         params.dt_bias,
                                         params.beta,
                                         params.threshold);
  // return npu::npu_fused_gdn_gating(params.A_log,
  //                                  params.a,
  //                                  params.b,
  //                                  params.dt_bias,
  //                                  params.beta,
  //                                  params.threshold);
#else
  NOT_IMPLEMENTED();
#endif
}

std::pair<torch::Tensor, torch::Tensor> fused_recurrent_gated_delta_rule(
    FusedRecurrentGatedDeltaRuleParams& params) {
#if defined(USE_NPU)
  return npu::npu_fused_recurrent_gated_delta_rule(
      params.q,
      params.k,
      params.v,
      params.g,
      params.beta,
      params.scale,
      params.initial_state,
      params.inplace_final_state,
      params.cu_seqlens,
      params.ssm_state_indices,
      params.num_accepted_tokens,
      params.use_qk_l2norm_in_kernel);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params) {
#if defined(USE_CUDA)
  auto out_2d = cuda::fp8_scaled_matmul(params.a,
                                        params.b,
                                        params.a_scale,
                                        params.b_scale,
                                        params.output_dtype,
                                        params.bias,
                                        params.output);

  // Auto reshape output if original input shape is provided
  if (params.input_shape.has_value()) {
    auto out_shape = params.input_shape.value();
    out_shape.back() = params.b.size(0);
    return out_2d.view(out_shape);
  }
  return out_2d;
#else
  LOG(FATAL) << "fp8_scaled_matmul is only supported on CUDA";
  return torch::Tensor();
#endif
}

void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params) {
#if defined(USE_CUDA)
  cuda::static_scaled_fp8_quant(params.output, params.input, params.scale);
#else
  LOG(FATAL) << "static_scaled_fp8_quant is only supported on CUDA";
#endif
}

// Fused RMSNorm + Static FP8 Quantization
torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params) {
#if defined(USE_CUDA)
  auto org_shape = params.input.sizes().vec();
  auto hidden_size = params.input.size(-1);

  // Flatten input to 2D. Use reshape to support non-contiguous tensors.
  auto input_2d = params.input.reshape({-1, hidden_size});

  torch::Tensor output =
      torch::empty({input_2d.size(0), hidden_size},
                   input_2d.options().dtype(torch::kFloat8_e4m3fn));

  // Call fused kernel
  cuda::rms_norm_static_fp8_quant(
      output, input_2d, params.weight, params.scale, params.epsilon);

  return output.reshape(org_shape);
#else
  LOG(FATAL) << "rms_norm_static_fp8_quant is only supported on CUDA";
  return torch::Tensor();
#endif
}

std::tuple<torch::Tensor, torch::Tensor> fused_add_rms_norm_static_fp8_quant(
    FusedAddRmsNormStaticFp8QuantParams& params) {
#if defined(USE_CUDA)
  auto org_shape = params.input.sizes().vec();
  auto hidden_size = params.input.size(-1);

  // Flatten tensors to 2D. Use reshape to support non-contiguous tensors.
  auto input_2d = params.input.reshape({-1, hidden_size});
  auto residual_2d = params.residual.reshape({-1, hidden_size});

  torch::Tensor output =
      torch::empty({input_2d.size(0), hidden_size},
                   input_2d.options().dtype(torch::kFloat8_e4m3fn));

  // Call fused kernel (residual is updated in-place)
  cuda::fused_add_rms_norm_static_fp8_quant(output,
                                            input_2d,
                                            residual_2d,
                                            params.weight,
                                            params.scale,
                                            params.epsilon);

  // Reshape outputs
  auto output_reshaped = output.reshape(org_shape);
  auto residual_reshaped = residual_2d.reshape(org_shape);

  return std::make_tuple(output_reshaped, residual_reshaped);
#else
  LOG(FATAL) << "fused_add_rms_norm_static_fp8_quant is only supported on CUDA";
  return std::make_tuple(torch::Tensor(), torch::Tensor());
#endif
}

torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params) {
#if defined(USE_NPU)
  if (params.conv_state_indices.has_value()) {
    CHECK(params.conv_state_indices.value().is_contiguous())
        << "causal_conv1d_update: conv_state_indices must be contiguous.";
  }
  return npu::npu_causal_conv1d_update_v2(params.x,
                                          params.conv_state,
                                          params.weight,
                                          params.activation,
                                          params.bias,
                                          params.conv_state_indices,
                                          params.query_start_loc,
                                          params.max_query_len,
                                          params.pad_slot_id,
                                          params.block_idx_last_scheduled_token,
                                          params.initial_state_idx,
                                          params.validate_data);

#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor gated_layer_norm(GatedLayerNormParams& params) {
#if defined(USE_NPU)
  return npu::layer_norm_fwd(params.x,
                             params.weight,
                             params.bias,
                             params.eps,
                             params.z,
                             params.group_size,
                             params.norm_before_gate,
                             params.is_rms_norm);
#elif defined(USE_MLU)
  return mlu::gated_layer_norm(params.x,
                               params.weight,
                               params.bias,
                               params.eps,
                               params.z,
                               params.group_size,
                               params.norm_before_gate);
#else
  NOT_IMPLEMENTED();
#endif
}

std::pair<torch::Tensor, torch::Tensor> partial_rotary_embedding(
    PartialRotaryEmbeddingParams& params) {
#if defined(USE_NPU)
  return npu::apply_npu_partial_rotary_embedding(params.positions,
                                                 params.query,
                                                 params.key,
                                                 params.head_size,
                                                 params.rotary_dim,
                                                 params.cos_sin_cache,
                                                 params.is_neox_style);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params) {
#if defined(USE_NPU)
  return npu::npu_fused_qkvzba_split_reshape_cat(params.mixed_qkvz,
                                                 params.mixed_ba,
                                                 params.num_heads_qk,
                                                 params.num_heads_v,
                                                 params.head_qk,
                                                 params.head_v);
#else
  NOT_IMPLEMENTED();
#endif
}

void gemma_rms_norm(GemmaRMSNormParams& params) {
#if defined(USE_NPU)
  npu::npu_gemma_rms_norm(
      params.x, params.gamma, params.epsilon, params.rstd_out, params.norm_out);
#elif defined(USE_MLU)
  mlu::gemma_rms_norm(params.x, params.gamma, params.epsilon, params.norm_out);
#else
  NOT_IMPLEMENTED();
#endif
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params) {
#if defined(USE_NPU)
  return npu::tilelang::split_qkv_rmsnorm_mrope(params.qkvg,
                                                params.q_weight,
                                                params.k_weight,
                                                params.cos_sin,
                                                params.gather_pattern,
                                                params.eps,
                                                params.num_q_heads,
                                                params.num_kv_heads,
                                                params.head_size);
#else
  NOT_IMPLEMENTED();
#endif
}

bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads,
                                                int64_t num_kv_heads,
                                                int64_t head_size) {
#if defined(USE_NPU)
  return npu::tilelang::has_split_qkv_rmsnorm_mrope_specialization(
      num_q_heads, num_kv_heads, head_size);
#else
  return false;
#endif
}

torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern(
    int64_t rope_dim,
    const std::vector<int64_t>& mrope_section,
    bool is_interleaved,
    const torch::Device& device) {
#if defined(USE_NPU)
  return npu::tilelang::build_split_qkv_rmsnorm_mrope_gather_pattern(
      rope_dim, mrope_section, is_interleaved, device);
#else
  NOT_IMPLEMENTED();
#endif
}

std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
    ChunkGatedDeltaRuleParams& params) {
#if defined(USE_NPU)
  return npu::npu_chunk_gated_delta_rule(params.q,
                                         params.k,
                                         params.v,
                                         params.g,
                                         params.beta,
                                         params.scale,
                                         params.initial_state,
                                         params.output_final_state,
                                         params.cu_seqlens,
                                         params.head_first,
                                         params.use_qk_l2norm_in_kernel);
#else
  NOT_IMPLEMENTED();
#endif
}

torch::Tensor recurrent_gated_delta_rule(
    const torch::Tensor& query,
    const torch::Tensor& key,
    const torch::Tensor& value,
    torch::Tensor& state,
    const std::optional<torch::Tensor>& beta,
    const std::optional<double> scale,
    const std::optional<torch::Tensor>& actual_seq_lengths,
    const std::optional<torch::Tensor>& ssm_state_indices,
    const std::optional<torch::Tensor>& num_accepted_tokens,
    const std::optional<torch::Tensor>& g,
    const std::optional<torch::Tensor>& gk) {
#if defined(USE_NPU)
  return npu::npu_recurrent_gated_delta_rule(query,
                                             key,
                                             value,
                                             state,
                                             beta,
                                             scale,
                                             actual_seq_lengths,
                                             ssm_state_indices,
                                             num_accepted_tokens,
                                             g,
                                             gk);
#else
  NOT_IMPLEMENTED();
#endif
}
}  // namespace xllm::kernel