// SPDX-License-Identifier: MIT

#pragma once

#include <torch/extension.h>

using namespace at;

void fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(Tensor& qkv,
                                                    Tensor& qw,
                                                    Tensor& kw,
                                                    Tensor& cos_sin,
                                                    Tensor& positions,
                                                    int64_t num_tokens,
                                                    int64_t num_heads_q,
                                                    int64_t num_heads_k,
                                                    int64_t num_heads_v,
                                                    int64_t head_size,
                                                    bool is_neox_style,
                                                    std::vector<int64_t> mrope_section_,
                                                    bool is_interleaved,
                                                    double eps,
                                                    Tensor& q_out,
                                                    Tensor& k_cache,
                                                    Tensor& v_cache,
                                                    Tensor& slot_mapping,
                                                    Tensor& per_tensor_k_scale,
                                                    Tensor& per_tensor_v_scale,
                                                    std::optional<Tensor> k_out,
                                                    std::optional<Tensor> v_out,
                                                    bool return_kv,
                                                    bool use_shuffle_layout,
                                                    int64_t block_size,
                                                    int64_t x,
                                                    int64_t rotary_dim = 0);
