import torch
import torch_mlu
import inspect
from ._utils import torchDtype2Str as _torchDtype2Str
from typing import Tuple, List, Dict, Optional, Union
from ._utils import get_custom_op_library_path as _get_custom_op_library_path
from ._utils import add_gen_case_decorator

torch.ops.load_library(_get_custom_op_library_path())

__all__ = [
    "active",
    "apply_rotary",
    "attention_project",
    "batch_matmul",
    "copy_blocks",
    "dequant_from_linear_cache",
    "dequant_from_paged_cache",
    "ffn",
    "flash_attention",
    "flash_attn_sq_mm_allreduce",
    "fused_layer_norm",
    "fused_moe",
    "fused_norm_attention_project",
    "fused_norm_residual_ffn",
    "fused_rms_norm",
    "fused_rope",
    "group_gemm",
    "matmul",
    "matmul_allreduce",
    "moe_active",
    "moe_cast_gating",
    "moe_combine_result",
    "moe_expand_input",
    "moe_gen_idx",
    "moe_quantize",
    "moe_softmax_topk",
    "offline_quant_to_linear_cache",
    "offline_quant_to_paged_cache",
    "per_token_smooth_quantize",
    "preload",
    "quant_to_linear_cache",
    "quant_to_paged_cache",
    "quantize",
    "reshape_linear_cache",
    "reshape_paged_cache",
    "single_query_cached_kv_attn",
    "single_query_mixed_cached_kv_attn",
    "smooth_quant_group_gemm",
    "smooth_quant_matmul",
    "smooth_quant_matmul_allreduce",
    "swap_blocks",
    "update_out_and_lse",
    "weight_only_quant_matmul",
]

def quantize(x: torch.Tensor,
             scale: torch.Tensor,
             zero: torch.Tensor = None
            ) -> torch.Tensor:
    """
    Apply quantization to the input tensor.

    Math:
        scaled = x * scale
        output = scaled.round().clamp(-128, 127).to(torch.int8)

    Args:
        x (torch.Tensor): The tensor to be quantized. Shape is (..., C).
        scale (torch.Tensor): The scale multipled to the input tensor. Shape is (C).
        zero (torch.Tensor): Not supported, must pass None.

    DataType:
        x: float, half or bfloat16
        scale: float
        output: int8

    Return:
        A int8-tensor with the same shape of x.
    """
    output = torch.empty_like(x, dtype=torch.int8, device=x.device)
    torch.ops.torch_mlu_ops.smooth_quant(x, scale, output, torch.Tensor(), zero,
                                                  None, None, None, 'per_token', False)
    return output

def per_token_smooth_quantize(x: torch.Tensor,
                              smooth: torch.Tensor,
                              zero: torch.Tensor = None,
                              token_count: torch.Tensor = None
                             ) -> Tuple[torch.Tensor]:
    """
    Apply per_token smooth quantization to the input tensor.

    Math:
        if token_count:
            input_list = input.split(token_count)
            group = token_count.size()
            result = []
            for i in range(group):
                result.append(input_list[i] * smooth[i])
            smoothed = concat(result, dim=0)
        else:
            smoothed = input * smooth
        max, _ = smoothed.abs().max(dim=-1, keepdim=True)
        scale = max.to(torch.float) / 127.0
        quanted = (smoothed / scale).round().clamp(-128, 127).to(torch.int8)
        output = (quanted, scale)

    Args:
        input (torch.Tensor): The tensor to be quantized. Shape is (..., C) or (token_count.sum(), C).
        smooth (torch.Tensor): The smooth scale multipled to the input tensor. Shape is (C) or (group, C).
        zero (torch.Tensor): Not supported, must pass None.
        token_count (torch.Tensor): The tensor to separate input. Shape is (group).

    DataType:
        x: float, half or bfloat16
        smooth: float
        token_count: int32
        output: (int8, float)

    Return:
        A tuple of two tensors.
        quantized_x: int8-tensor with the same shape of x.
        scale: float-tensor with the shape of x.shape[0:-1].
    """
    output = torch.empty(x.size(), dtype=torch.int8, device=x.device)
    output_scale = torch.empty(x.size()[:-1], dtype=smooth.dtype, device=smooth.device)
    torch.ops.torch_mlu_ops.smooth_quant(x, smooth, output, output_scale, zero,
                                         token_count, None, None, 'per_token', True)
    return (output, output_scale)

def weight_only_quant_matmul(a: torch.Tensor,
                             b: torch.Tensor,
                             scale: torch.Tensor,
                             zero: torch.Tensor = None,
                             bias: torch.Tensor = None,
                             c: torch.Tensor = None,
                             act_mode: str = "none",
                             quant_bit_size: int = 8,
                             alpha: float = 1.0,
                             beta: float = 1.0,
                             use_hp_active: bool = False):
    """
    Perform matrix multiplication on tensor a and quantized tensor b.
    Math:
        ab = matmul(a, b) * scale + bias
        acted = active(ab)
        output = alpha * acted + beta * c

    Args:
        a (torch.Tensor): Shape is (M, K).
        b (torch.Tensor): If quant_bit_size = 8, shape is (N, K).
                          If quant_bit_size = 4, shape is (N, K//2).
        scale (torch.Tensor): If use groupwise quantization, shape must be (N, group_num), data type must be
                              the same as a; otherwise shape must be (N), data tyte must be float.
        zero (torch.Tensor): Not supported, must pass None.
        bias (torch.Tensor): Shape is (N).
        c (torch.Tensor): Shape is (M, N).
        act_mode (str): Choose the activation algorithm, must be 'silu', 'gelu' or 'none'. If use groupwise
                        quantization, act_mode must be 'none'.
        quant_bit_size (int): The data format of b.
        alpha (float): coefficient of acted.
        beta (float): coefficient of c.
        use_hp_active (bool): Describing the algorithm that used in the implementation of the activation function.
             When the value is true, use the high-precision algorithm, otherwise use the fastest algorithm of activation.

    Type:
        a: float, half, bfloat16
        b: int8
        scale: float or same as a
        zero: float
        bias: same type as a
        c: same type as a
        output: same type as a

    Return:
        A tensor with the shape of (M, N).
    """
    is_group_quant = scale.dim() == 2
    b_scale = scale if is_group_quant else None
    gemm_out_scale = None if is_group_quant else scale
    b_quant_layout = 'quantize_group_wise' if is_group_quant else 'quantize_per_channel'
    return quant_matmul(a, None, None, b, b_scale, None, bias, c, None, None,
                        gemm_out_scale, zero, None, "weight_only", "quantize_none", b_quant_layout,
                        quant_bit_size, act_mode, use_hp_active, 1.0, alpha, beta, False, True)

def smooth_quant_matmul(a: torch.Tensor,
                        a_scale: Optional[torch.Tensor],
                        b: torch.Tensor,
                        b_scale: torch.Tensor,
                        dtype: torch.dtype,
                        bias: torch.Tensor = None,
                        c: torch.Tensor = None,
                        act_mode: str = "none",
                        alpha: float = 1.0,
                        beta: float = 1.0,
                        use_hp_active: bool = False):
    """
    Perform smooth quantized matrix multiplication on tensor a and b.

    Math:
        if a_scale is not None:
            ab = matmul(a, b) * a_scale.outer(b_scale) + bias
        else:
            ab = matmul(a, b) * b_scale + bias
        acted = active(ab)
        output = alpha * acted + beta * c

    Args:
        a (torch.Tensor): Shape is (M, K).
        a_scale (torch.Tensor): The per_token scale of a, shape is (M). If None, quant_layout is "quantize_per_tensor", and scale should be converted into b_scale.
        b (torch.Tensor): Shape is (N, K).
        b_scale (torch.Tensor): The per_channel scale of b, shape is (N).
        dtype (torch.dtype): Specify the data type of output, must be torch.float, torch.half or torch.bfloat16.
        bias (torch.Tensor): Shape is (N).
        c (torch.Tensor): Shape is (M, N).
        act_mode (str): Choose the activation algorithm, must be 'silu', 'gelu' or 'none'.
        alpha (float): coefficient of acted.
        beta (float): coefficient of c.
        use_hp_active (bool): Describing the algorithm that used in the implementation of the activation function.
             When the value is true, use the high-precision algorithm, otherwise use the fastest algorithm of activation.

    Type:
        a: int8
        a_scale: float
        b: int8
        b_scale: float
        bias: same as output
        c: same as output
        output: specified by dtype

    Return:
        A tensor with the shape of (M, N).
    """
    is_pertoken = a_scale is not None
    assert dtype in (torch.half, torch.float, torch.bfloat16), "dtype must be half, float or bfloat16, got: {}.".format(dtype)
    return quant_matmul(a, a_scale, None, b, b_scale, None, bias, c, None, None, b_scale, None,
                        _torchDtype2Str(dtype), "smooth_quant",
                        "quantize_per_token" if is_pertoken else "quantize_per_tensor",
                        "quantize_per_channel",
                        8, act_mode, use_hp_active, 1.0, alpha, beta, False, True)

def smooth_quant_matmul_allreduce(cncl_comm,
                                  a: torch.Tensor,
                                  a_scale: torch.Tensor,
                                  b: torch.Tensor,
                                  b_scale: torch.Tensor,
                                  dtype: torch.dtype,
                                  bias: torch.Tensor = None,
                                  c: torch.Tensor = None,
                                  alpha: float = 1.0,
                                  beta: float = 1.0,
                                  block_m: int = 0):
    """
    Parallelly apply smooth quantized matrix multiplication with TP-allreduce.

    Math:
        if a_scale is not None:
            ab = matmul(a, b) * a_scale.outer(b_scale) + bias
        else:
            ab = matmul(a, b) * b_scale + bias
        acted = active(ab)
        output = alpha * acted + beta * c
        output_allreduce = allreduce(output)

    Args:
        cncl_comm (int): The handle of CnclComm.
        a (torch.Tensor): Shape is (M, K).
        a_scale (torch.Tensor): The per_token scale of a, shape is (M)
        b (torch.Tensor): Shape is (N, K).
        b_scale (torch.Tensor): The per_channel scale of b, shape is (N).
        dtype (torch.dtype): Specify the data type of output, must be torch.float, torch.half or torch.bfloat16.
        bias (torch.Tensor): Shape is (N).
        c (torch.Tensor): Shape is (M, N).
        alpha (float): coefficient of acted.
        beta (float): coefficient of c.
        block_m (int): If the value is greater than 0, the specified value will be used; if the value is less than 0, the recommended value will be used.

    Type:
        a: int8
        a_scale: float
        b: int8
        b_scale: float
        bias: same as output
        c: same as output
        output: specified by dtype

    Return:
        A tensor with the shape of (M, N).
    """
    is_pertoken = a_scale is not None
    assert dtype in (torch.half, torch.float, torch.bfloat16), "dtype must be half, float or bfloat16, got: {}.".format(dtype)
    return quant_matmul_allreduce(cncl_comm, a, a_scale, None, b, b_scale,
            None, bias, c, None, None, b_scale, None,
            _torchDtype2Str(dtype), "smooth_quant",
            "quantize_per_token" if is_pertoken else "quantize_per_tensor",
            "quantize_per_channel",
            8, alpha, beta, False, True, block_m)

def fused_layer_norm(x: torch.Tensor,
                     residual: Optional[torch.Tensor],
                     gamma: Optional[torch.Tensor],
                     beta: Optional[torch.Tensor],
                     bias: Optional[torch.Tensor],
                     eps: float,
                     store_output_before_norm: bool,
                     quant_scale: torch.Tensor = None,
                     out: torch.Tensor = None,
                     dynamic_quant: bool = False):
    """
    Apply layernorm to the input tensor.
    Please refer to https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html#torch.nn.LayerNorm

    Math:
        if bias is not None:
            x1 = x + bias
        if residual is not None:
            x1 += residual
        normed = layernorm(x1, gamma, beta, eps)
        if dynamic_quant:
            max, _ = (normed * quant_scale).abs().max(dim=-1, keepdim=True)
            smooth_quant_scale = max.to(torch.float) / int_max
            normed = (normed / smooth_quant_scale).round().to(int8)
        elif quant_scale is not None:
            normed = (normed * quant_scale).round().clamp(-128, 127).to(torch.int8)
        if store_output_before_norm == True:
            output = (normed, x1)
        else:
            output = normed

    Args:
        x (torch.Tensor): The input tensor to be normalized. Shape is (p1, p2, ..., pn, C).
        residual (torch.Tensor): The tensor add to x. Shape is the same as x.
        gamma (torch.Tensor): The weight of layernorm. Shape is (C).
        beta (torch.Tensor): The bias of layernorm. Shape is (C).
        eps (float): The eps of layernorm.
        store_output_before_norm (bool): If return the tensor before applying layernorm.
        quant_scale (torch.Tensor): The smooth_scale of output if dynamic_quant, else it is the scale of quantization. Shape is (C).
        out (torch.Tensor): The output tensor of normalizes. Shape is the same as x.
        dynamic_quant(bool): Flag quant_mode of output, dynamic quant or not.
        smooth_quant_scale(torch.Tensor): The output scale tensor. Shape is [p1, p2, ..., pn].

    Type:
        x: float, half, bfloat16
        residual: same as x
        beta: same as x
        gamma: same as x
        bias: same as x
        quant_scale: float
        output: same as x, int8
        smooth_quant_scale: float

    Return:
        Support inplace output.
        If store_output_before_norm is True, return Two tensors before and after applying layernorm,
        otherwise return a tensor after applying layernorm and quantization.
    """
    if out is None:
        if quant_scale is not None:
            out = torch.empty(*x.shape, dtype=torch.int8, device=x.device)
        else:
            out = torch.empty(*x.shape, dtype=x.dtype, device=x.device)
    else:
        assert store_output_before_norm is False and quant_scale is None, "store_output_before_norm and quant_out are not support when inplace out."
    residual_out = torch.empty(*x.shape, dtype=x.dtype, device=x.device) if store_output_before_norm else None
    smooth_quant_scale = torch.empty(x.shape[:-1], dtype=torch.float, device=x.device) if dynamic_quant else None
    torch.ops.torch_mlu_ops.fused_layernorm(x, out, residual, gamma, beta, bias, quant_scale, residual_out, smooth_quant_scale, "layernorm", eps, store_output_before_norm, dynamic_quant)
    output = [out]
    if store_output_before_norm:
        output.append(residual_out)
    if dynamic_quant:
        output.append(smooth_quant_scale)
    return output[0] if len(output) == 1 else tuple(output)


def fused_rms_norm(x: torch.Tensor,
                   residual: Optional[torch.Tensor],
                   gamma: Optional[torch.Tensor],
                   beta: Optional[torch.Tensor],
                   bias: Optional[torch.Tensor],
                   eps: float,
                   store_output_before_norm: bool,
                   quant_scale: torch.Tensor = None,
                   out: torch.Tensor = None,
                   dynamic_quant: bool = False):
    """
    Apply rmsnorm to the input tensor.
    Paramter limitations the same as "fused_layer_norm", except 'beta' should be None.
    """
    if out is None:
        if quant_scale is not None:
            out = torch.empty(*x.shape, dtype=torch.int8, device=x.device)
        else:
            out = torch.empty(*x.shape, dtype=x.dtype, device=x.device)
    else:
        assert store_output_before_norm is False and quant_scale is None, "store_output_before_norm and quant_out are not support when inplace out."
    residual_out = torch.empty(*x.shape, dtype=x.dtype, device=x.device) if store_output_before_norm else None
    smooth_quant_scale = torch.empty(x.shape[:-1], dtype=torch.float, device = x.device) if dynamic_quant else None
    torch.ops.torch_mlu_ops.fused_layernorm(x, out, residual, gamma, beta, bias, quant_scale, residual_out, smooth_quant_scale, "rmsnorm", eps, store_output_before_norm, dynamic_quant)
    output = [out]
    if store_output_before_norm:
        output.append(residual_out)
    if dynamic_quant:
        output.append(smooth_quant_scale)
    return output[0] if len(output) == 1 else tuple(output)

def attention_project(input: torch.Tensor,
                      weight: torch.Tensor,
                      bias: torch.Tensor = None,
                      residual: torch.Tensor = None,
                      alpha: float = 1.0,
                      beta: float = 1.0
                      ) -> torch.Tensor:
    """
    Perform attention project operation which including adding bias and residual.

    Math:
        proj = matmul(input, weight)
        if bias is not None:
            proj += bias
        if residual is not None:
            proj = alpha * proj + beta * residual
        output = proj

    Args:
        input (torch.Tensor): The input tensor. Shape is (batch, seq, input_size) or (total_seq, input_size).
        weight (torch.Tensor): The weight tensor. Shape is (hidden_size, input_size).
        bias (torch.Tensor): The bias tensor. Shape is (hidden_size).
        residual (torch.Tensor): The residual tensor. Shape is the same as input.
        alpha (float): The coefficient of proj.
        beta (float): The coefficient of residual.

    Type:
        input: float, half, bfloat16
        weight: same as input
        bias: same as input
        residual: same as input
        output: same as input

    Return:
        A tensor with the shape of (batch, seq, hidden_size) or (total_seq, hidden_size).
    """
    output = torch.ops.torch_mlu_ops.attention_project(input, weight, bias, None, None, None,
                                       None, None, None, residual, "nthc", 1,
                                       1e-5, alpha, beta, False)
    return output[0]

def fused_norm_attention_project(input: torch.Tensor,
                                 q_weight: torch.Tensor,
                                 q_bias: Optional[torch.Tensor],
                                 k_weight: Optional[torch.Tensor],
                                 k_bias: Optional[torch.Tensor],
                                 v_weight: Optional[torch.Tensor],
                                 v_bias: Optional[torch.Tensor],
                                 norm_weight: Optional[torch.Tensor],
                                 norm_bias: Optional[torch.Tensor],
                                 eps: float,
                                 out_layout: str = "nthc",
                                 head_size: int = 1,
                                 norm_out: bool = False):
    """
    Perform attention project operation which including layernorm, adding bias/residual and permuting.

    Math:
        normed = layernorm(input, norm_weight, norm_bias, eps)
        proj_q = matmul(normed, q_weight) + q_bias
        if k_weight is not None:
            proj_k = matmul(normed, k_weight) + k_bias
        if v_weight is not None:
            proj_v = matmul(normed, v_weight) + v_bias
        if out_layout is 'nhtc':
            c = proj_q.reshape(batch, seq, hidden_size_q // head_size, head_size).transpose(1, 2)
            proj_k = proj_k.reshape(batch, seq, hidden_size_k // head_size, head_size).transpose(1, 2)
            proj_v = proj_v.reshape(batch, seq, hidden_size_v // head_size, head_size).transpose(1, 2)
        outputs = [proj_q]
        if k_weight is not None:
            outputs.append(proj_k)
        if v_weight is not None:
            outputs.append(proj_v)
        if norm_out:
            outputs.append(normed)
        return tuple(outputs)
    Args:
        input (torch.Tensor): The input tensor. Shape is (batch, seq, input_size) or (total_seq, input_size).
        q_weight (torch.Tensor): The q_weight tensor. Shape is (hidden_size_q, input_size).
        q_bias (torch.Tensor): The q_bias tensor. Shape is (hidden_size_q).
        k_weight (torch.Tensor): The k_weight tensor. Shape is (hidden_size_k, input_size).
        k_bias (torch.Tensor): The k_bias tensor. Shape is (hidden_size_k).
        v_weight (torch.Tensor): The v_weight tensor. Shape is (hidden_size_v, input_size).
        v_bias (torch.Tensor): The v_bias tensor. Shape is (hidden_size_v).
        norm_weight (torch.Tensor): The weight of layernorm. Shape is (input_size).
        norm_bias (torch.Tensor): The bias of layernorm. Shape is (input_size).
        eps (float): The eps of layernorm.
        out_layout (str): The layout of output tensor. Must be 'nthc' or 'nhtc'.
        head_size (int): The last dimension of output when out_layout is 'nhtc'.
        norm_out (bool): If return the layernorm result.

    Type:
        input: float, half, bfloat16
        q_weight: same as input
        q_bias: same as input
        k_weight: same as input
        k_bias: same as input
        v_weight: same as input
        v_bias: same as input
        norm_weight: same as input
        norm_bias: same as input
        output: same as input

    Return:
        If len(outputs) > 1, return a tuple of tensors, otherwise return a tensor.

    """
    output = torch.ops.torch_mlu_ops.attention_project(input, q_weight, q_bias, k_weight, k_bias, v_weight, v_bias,
                            norm_weight, norm_bias, None, out_layout, head_size, eps, 1.0,
                            1.0, norm_out)
    return tuple(output) if len(output) > 1 else output[0]

def ffn(input: torch.Tensor,
        up_fc_weight: torch.Tensor,
        up_fc_bias: Optional[torch.Tensor],
        down_proj_weight: torch.Tensor,
        down_proj_bias: Optional[torch.Tensor],
        gate_up_proj_weight: Optional[torch.Tensor],
        gate_up_proj_bias: Optional[torch.Tensor],
        act_mode: str
        ) -> torch.Tensor:
    """
    Perform feed forward operation.

    Math:
        up = matmul(input, up_fc_weight)
        if up_fc_bias is not None:
            up += up_fc_bias
        act = active(up)
        if gate_up_proj_weight is not None:
            gate = matmul(input, gate_up_proj_weight)
            if gate_up_proj_bias is not None:
                gate += gate_up_proj_bias
            act *= gate
        down = matmul(act, down_proj_weight)
        if gate_up_proj_bias is not None:
            down += gate_up_proj_bias
        output = down

    Args:
        input (torch.Tensor): The input tensor. Shape is (batch, seq, hidden_size) or (total_seq, hidden_size).
        up_fc_weight (torch.Tensor): The weight of upper matmul. Shape is (inner_size, hidden_size).
        up_fc_bias (torch.Tensor): The bias of upper matmul. Shape is (inner_size).
        down_proj_weight (torch.Tensor): The weight of lower matmul. Shape is (hidden_size, inner_size).
        down_proj_bias (torch.Tensor): The bias of lower matmul. Shape is (hidden_size).
        gate_up_proj_weight (torch.Tensor): The weight of gate matmul. Shape is (inner_size, hidden_size).
        gate_up_proj_bias (torch.Tensor): The bias of gate matmul. Shape is (inner_size).
        act_mode (str): The activation mode, must be 'silu', 'gelu', 'relu' or 'none'.

    Type:
        input: float, half, bfloat16.
        up_fc_weight: same as input.
        up_fc_bias: same as input.
        down_proj_weight: same as input.
        down_proj_bias: same as input.
        gate_up_proj_weight: same as input.
        gate_up_proj_bias: same as input.
        output: same as input.

    Return:
        An output tensor.
    """
    return torch.ops.torch_mlu_ops.ffn(input, up_fc_weight, up_fc_bias, down_proj_weight, down_proj_bias,
               gate_up_proj_weight, gate_up_proj_bias, None, None, act_mode, "none",
               1e-5, 1., 1.)

def fused_norm_residual_ffn(input: torch.Tensor,
                            up_fc_weight: torch.Tensor,
                            up_fc_bias: Optional[torch.Tensor],
                            down_proj_weight: torch.Tensor,
                            down_proj_bias: Optional[torch.Tensor],
                            gate_up_proj_weight: Optional[torch.Tensor],
                            gate_up_proj_bias: Optional[torch.Tensor],
                            layernorm_weight: Optional[torch.Tensor],
                            layernorm_bias: Optional[torch.Tensor],
                            eps: float,
                            act_mode: str,
                            residual_is: str = "input",
                            alpha: float = 1.0,
                            beta: float = 1.0
                            ) -> torch.Tensor:
    """
    Perform feed forward operation.

    Math:
        residual = input
        normed = input
        if layernorm_weight is not None:
            normed = layernorm(input, layernorm_weight, layernorm_bias, eps)
        if residual_is == 'normed_input':
            residual = normed
        up = matmul(normed, up_fc_weight)
        if up_fc_bias is not None:
            up += up_fc_bias
        act = active(up)
        if gate_up_proj_weight is not None:
            gate = matmul(input, gate_up_proj_weight)
            if gate_up_proj_bias is not None:
                gate += gate_up_proj_bias
            act *= gate
        down = matmul(act, down_proj_weight)
        if gate_up_proj_bias is not None:
            down += gate_up_proj_bias
        if residual_is != 'none':
            down = alpha * down + beta * residual
        output = down

    Args:
        input (torch.Tensor): The input tensor. Shape is (batch, seq, hidden_size) or (total_seq, hidden_size).
        up_fc_weight (torch.Tensor): The weight of upper matmul. Shape is (inner_size, hidden_size).
        up_fc_bias (torch.Tensor): The bias of upper matmul. Shape is (inner_size).
        down_proj_weight (torch.Tensor): The weight of lower matmul. Shape is (hidden_size, inner_size).
        down_proj_bias (torch.Tensor): The bias of lower matmul. Shape is (hidden_size).
        gate_up_proj_weight (torch.Tensor): The weight of gate matmul. Shape is (inner_size, hidden_size).
        gate_up_proj_bias (torch.Tensor): The bias of gate matmul. Shape is (inner_size).
        layernorm_weight (torch.Tensor): The weight of layernorm.
        layernorm_bias (torch.Tensor): The bias of layernorm.
        eps (float): The eps of layernorm.
        act_mode (str): The activation mode, must be 'silu', 'gelu', 'relu' or 'none'.
        residual_is (str): A string controlling which tensor is residual, must be 'input', 'normed_input' or 'none'.
        alpha (float): The coefficient of lower matmul result.
        bete (float): The coefficient of residual.

    Type:
        input: float, half, bfloat16.
        up_fc_weight: same as input.
        up_fc_bias: same as input.
        down_proj_weight: same as input.
        down_proj_bias: same as input.
        gate_up_proj_weight: same as input.
        gate_up_proj_bias: same as input.
        layernorm_weight: same as input.
        layernorm_bias: same as input.
        output: same as input.

    Return:
        An output tensor.
    """
    return torch.ops.torch_mlu_ops.ffn(input, up_fc_weight, up_fc_bias, down_proj_weight, down_proj_bias,
                gate_up_proj_weight, gate_up_proj_bias, layernorm_weight, layernorm_bias, act_mode,
                residual_is, eps, alpha, beta)

def flash_attention(q: torch.Tensor,
                    k: torch.Tensor,
                    v: torch.Tensor,
                    out: Optional[torch.Tensor],
                    cu_seq_lens_q: Optional[torch.Tensor],
                    cu_seq_lens_kv: Optional[torch.Tensor],
                    alibi_slope: Optional[torch.Tensor],
                    attn_bias: Optional[torch.Tensor],
                    max_seq_len_q: int,
                    max_seq_len_kv: int,
                    softmax_scale: float,
                    is_causal: bool,
                    window_size_left: int = -1,
                    window_size_right: int = -1,
                    compute_dtype: torch.dtype = torch.float,
                    return_lse: bool = False,
                    block_tables: Optional[torch.Tensor] = None,
                    k_cache_quant_scale: Optional[torch.Tensor] = None,
                    v_cache_quant_scale: Optional[torch.Tensor] = None
                    ):
    """
    Apply attention operation on q, k and v.

    Math:
        qk = bmm(q, k) * softmax_scale
        if alibi_slope is not None:
            qk += create_alibi(alibi_slope)
        if attn_bias is not None:
            qk += attn_bias
        if is_causal:
            qk = tril_mask(qk)
        if return_lse:
            lse = logsumexp(qk, dim=-1)
        attn = softmax(qk, dim=-1)
        attn_out = bmm(attn, v)
        output = (attn_out, lse) if return_lse else attn_out

    Args:
        q (torch.Tensor): The query tensor. Shape is (batch, seq_q, head_num_q, head_size_qk) or (total_seq_q, head_num_q, head_size_qk).
        k (torch.Tensor): The key tensor. If block_tables is None, the shape of k is (batch, seq_kv, head_num_kv, head_size_qk) or (total_seq_kv, head_num_kv, head_size_qk).
                                          If block_tables is not None, if use paged attention, shape is (total_blocks, head_num_kv, block_size, head_size_qk),
                                            where the total_blocks = max_batch * (memory_cache_len / block_size). If not use paged attention, shape is (max_batch, head_num_kv, memory_cache_len, head_size_qk).
        v (torch.Tensor): The value tensor. Shape is the same as k, except, head_size_v could be different with head_size_qk.
        out (torch.Tensor): The output tensor. Shape is the same as q.
        cu_seq_lens_q (torch.Tensor): The cusum of seq_q if q is 3-D tensor. Shape is (batch+1).
        cu_seq_lens_kv (torch.Tensor): The cusum of seq_kv if k/v is 3-D tensor. Shape is (batch+1).
        alibi_slope (torch.Tensor): The alibi_slope to generate alibi position embedding. Shape is (head_num_q).
        attn_bias (torch.Tensor): The bias added to qk. Shape is (batch, head_num_q, max_seq_len_q, max_seq_len_kv).
        max_seq_len_q (int): The maximum value of seq_q.
        max_seq_len_kv (int): The maximum value of seq_kv.
        softmax_scale (float): The scale factor multiplied to qk.
        is_causal (bool): Controlling if use causal mask.
        window_size_left (int): left window size, set -1 if it is unlimited. The maximum lengths of key and value involved in the attention calculation is windows_size_left + 1.
        window_size_right (int): right window size, set -1 if it is unlimited. The maximum lengths of key and value involved in the attention calculation is windows_size_right + 1
        compute_dtype (torch.dtype): The dtype used to calculate softmax.
        return_lse (bool): Controlling if return lse tensor.
        k_cache_quant_scale (torch.Tensor): Reserved parameter, the quantized scale of k_cache, must be None.
        v_cache_quant_scale (torch.Tensor): Reserved parameter, The quantized scale of v_cache, must be None.
        block_tables (torch.Tensor): The tensor recording the k/v_cache's positions of each batch.
            If use paged attention, shape is  (batch, max_block_num), where max_block_num = memory_cache_len / block_size.
            If not use paged attention, shape is (batch, 1).

    Type:
        q: float, half, bfloat16.
        k: same as q.
        v: same as q.
        cu_seq_lens_q: int32.
        cu_seq_lens_kv: int32.
        alibi_slope: float.
        attn_bias: same as q.
        attn_out : same as q.
        lse: float.

    Return:
        If return_lse is True, return a tuple of (attn_out, lse), else return attn_out.
    """
    tmo_out = out
    if out is None:
       tmo_out = torch.empty(q.size()[:-1] + (v.size()[-1],), dtype=q.dtype, device=q.device)
    out_lse = None
    if return_lse:
        is_pack = cu_seq_lens_q is not None
        batch = cu_seq_lens_q.shape[0] - 1 if is_pack else q.shape[0]
        head_num = q.shape[1] if is_pack else q.shape[2]
        out_lse = torch.empty((batch, head_num, max_seq_len_q), dtype=torch.float64, device=q.device)
    torch.ops.torch_mlu_ops.flash_attention(q, k, v,
                                            tmo_out, out_lse,
                                            cu_seq_lens_q, cu_seq_lens_kv, alibi_slope,
                                            attn_bias, k_cache_quant_scale, v_cache_quant_scale,
                                            block_tables, max_seq_len_q, max_seq_len_kv,
                                            softmax_scale, is_causal,
                                            window_size_left, window_size_right,
                                            _torchDtype2Str(compute_dtype), return_lse)
    return (tmo_out, out_lse) if return_lse else tmo_out


def single_query_cached_kv_attn(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
    out: torch.Tensor,
    block_tables: torch.Tensor,
    context_lens: torch.Tensor,
    k_cache_quant_scale: Optional[torch.Tensor],
    v_cache_quant_scale: Optional[torch.Tensor],
    alibi_slopes: Optional[torch.Tensor],
    max_contxt_len: int,
    windows_size_left: int,
    windows_size_right: int,
    softmax_scale: float,
    return_lse: bool = False,
    kv_cache_quant_bit_size = -1
) -> torch.Tensor:
    """
    Apply attention operation on query and cached key/value.

    Math:
        out_list = []
        lse_list = []
        for i in range(q.shape[0]):
            key = collect(k_cache, block_tables[i], context_lens[i])
            value = collect(v_cache, block_tables[i], context_lens[i])
            qk = bmm(q[i], key) * softmax_scale
            if alibi_slope is not None:
                alpha = qk + create_alibi(alibi_slope)
            attn = softmax(alpha, dim=-1)
            attn_out = bmm(attn, value)
            lse = logsumexp(qk, dim = -1)
            out_list.append(attn_out)
            lse_list.append(lse)
        output = concat(out_list, dim = 0)
        output_lse = concat(lse_list, dim = 0)

    Args:
        q (torch.Tensor): The query tensor, shape is (batch, seq_q, head_num_q, head_size_qk).
        k_cache (torch.Tensor): The key cache tensor.
            If use paged attention, shape is (total_blocks, head_num_kv, block_size, head_size_qk), where the total_blocks = max_batch * (memory_cache_len / block_size).
            If not use paged attention, shape is (max_batch, head_num_kv, memory_cache_len, head_size_qk).
        v_cache (torch.Tensor): The value cache tensor. shape is the same as k_cache, except, head_size_v could be different with head_size_qk.
        out (torch.Tensor): The output tensor, shape is the same as q
        block_tables (torch.Tensor): The tensor recording the k/v_cache's positions of each batch.
            If use paged attention, shape is (batch, max_block_num), where max_block_num = max_cache_len / block_size.
            If not use paged attention, shape is (batch, 1).
        context_lens (torch.Tensor): The tensor recording the context lengths of each batch, shape is (batch).
        k_cache_quant_scale (torch.Tensor): The quantized scale of k_cache.
                            When quantize mode is per_token, shape is
                            paged:(total_blocks, head_num_kv, block_size), linear: (max_batch, head_num_kv, memory_cache_len).
                            Or, paged:(total_blocks, head_num_kv, block_size, 1), linear: (max_batch, head_num_kv, memory_cache_len, 1).
                            when quantize mode is per_channel, shape is (head_num_kv, head_size)
        v_cache_quant_scale (torch.Tensor): The quantized scale of v_cache, shape is the same as k_cache_quant_scale.
        alibi_slopes (torch.Tensor): The alibi_slope to generate alibi position embedding. Shape is (head_num_q).
        max_contxt_len (int): The maximum value of context_lens.
                              If max_contxt_len <= 0, indicating an uncertain maximum value of context_lens.
                              If 0 < max_contxt_len < max(context_lens), the correctness is not guaranteed.
        windows_size_left (int): left window size, set -1 if it is unlimited. The maximum lengths of key and value involved in the attention calculation is windows_size_left + 1.
        windows_size_right (int): Reserved parameter, set -1.
        softmax_scale (float): The scale factor multiplied to qk.
        return_lse (bool): Controlling if return lse tensor.
        kv_cache_quant_bit_size (int): The data format of kv_cache, 4 for int4, 8 for int8, and -1 for the same as query.

    Type:
        q: float, half, bfloat16.
        k_cache: float, half, bfloat16, int8, int8(pairwise-int4).
        v_cache: same as k_cache.
        block_tables: int32.
        context_lens: int32.
        k_cache_quant_scale: float.
        v_cache_quant_scale: float.
        alibi_slopes: float.

    Return:
        If return_lse is True, return a tuple of (output, lse), else return output.
    """
    tmo_out = torch.empty((q.shape[0], q.shape[1], q.shape[2], v_cache.shape[-1]), dtype=q.dtype, device=q.device) if out is None else out
    out_lse = None
    if return_lse:
        batch = q.shape[0]
        seq_q = q.shape[1]
        head_num = q.shape[2]
        out_lse = torch.empty((batch, head_num, seq_q), dtype=torch.float, device=q.device)
    torch.ops.torch_mlu_ops.single_query_cached_kv_attn(
            q, k_cache, v_cache, tmo_out, block_tables, context_lens, out_lse,
            k_cache_quant_scale, v_cache_quant_scale, alibi_slopes, max_contxt_len,
            windows_size_left, windows_size_right, softmax_scale, return_lse,
            kv_cache_quant_bit_size)
    return (tmo_out, out_lse) if return_lse else tmo_out

def single_query_mixed_cached_kv_attn(
    q: torch.Tensor,
    k_cache_lp: torch.Tensor,
    v_cache_lp: torch.Tensor,
    k_cache_hp: torch.Tensor,
    v_cache_hp: torch.Tensor,
    out: torch.Tensor,
    block_tables_lp: torch.Tensor,
    block_tables_hp: torch.Tensor,
    context_lens_lp: torch.Tensor,
    context_lens_hp: torch.Tensor,
    k_cache_quant_scale_lp: Optional[torch.Tensor],
    v_cache_quant_scale_lp: Optional[torch.Tensor],
    k_cache_quant_scale_hp: Optional[torch.Tensor],
    v_cache_quant_scale_hp: Optional[torch.Tensor],
    alibi_slopes: Optional[torch.Tensor],
    max_contxt_len_lp: int,
    max_contxt_len_hp: int,
    softmax_scale: float,
    return_lse: bool = False,
    kv_cache_quant_bit_size_lp = -1,
    kv_cache_quant_bit_size_hp = -1,
) -> Union[Tuple[torch.Tensor], torch.Tensor]:
    """
    Apply attention operation on query and cached key/value.

    Math:
        lse1, attn_out1 = single_query_cached_kv_attn(query, key_cache_lp, value_cache_lp, ...)
        lse2, attn_out2 = single_query_cached_kv_attn(query, key_cache_hp, value_cache_hp, ...)
        output = update_out_and_lse(lse1, attn_out1, lse2, attn_out2)

    Args:
        q (torch.Tensor): The query tensor, shape is (batch, seq_q, head_num_q, head_size). seq_q must be 1 now.
        k_cache_lp/k_cache_hp (torch.Tensor): The key cache tensor.
            If use paged attention, shape is (total_blocks, head_num_kv, block_size, head_size), where the total_blocks = max_batch * (memory_cache_len / block_size).
            If not use paged attention, shape is (max_batch, head_num_kv, memory_cache_len, head_size).
        v_cache_lp/vcache_hp (torch.Tensor): The value cache tensor, shape is the same as k_cache, just, head_size_v could be different with head_size_qk.
        out (torch.Tensor): The output tensor, shape is the same as q
        block_tables_lp/block_tables_hp (torch.Tensor): The tensor recording the k/v_cache's positions of each batch.
            If use paged attention, shape is (batch, max_block_num), where max_block_num = memory_cache_len / block_size.
            If not use paged attention, shape is (batch, 1).
        context_lens_lp/context_lens_hp (torch.Tensor): The tensor recording the context lengths of each batch, shape is (batch).
        k_cache_quant_scale_lp/k_cache_quant_scale_hp (torch.Tensor): The quantized scale of k_cache.
                            When quantize mode is per_token, shape is
                            paged:(total_blocks, head_num_kv, block_size, 1), linear: (max_batch, head_num_kv, memory_cache_len, 1).
                            when quantize mode is per_channel, shape is (head_num_kv, head_size)
        v_cache_quant_scale_lp/v_cache_quant_scale_hp (torch.Tensor): The quantized scale of v_cache, shape is the same as k_cache_quant_scale.
        alibi_slopes (torch.Tensor): The alibi_slope to generate alibi position embedding. Shape is (head_num_q).
        max_contxt_len_lp/max_contxt_len_hp (int): The maximum value of context_lens.
                              If max_contxt_len <= 0, indicating an uncertain maximum value of context_lens.
                              If 0 < max_contxt_len < max(context_lens), the correctness is not guaranteed.
        softmax_scale (float): The scale factor multiplied to qk.
        return_lse (bool): Controlling if return lse tensor.
        kv_cache_quant_bit_size_lp/kv_cache_quant_bit_size_hp (int): The data format of kv_cache, -1 for float-point.

    Type:
        q: float, half, bfloat16.
        k_cache_lp/k_cache_hp: float, half, bfloat16, int8, int8(pairwise-int4).
        v_cache_lp/v_cache_hp: same as k_cache.
        block_tables_lp/block_tables_hp: int32.
        context_lens: int32.
        k_cache_quant_scale_lp/k_cache_quant_scale_hp: float.
        v_cache_quant_scale_lp/v_cache_quant_scale_hp: float.
        alibi_slopes: float.

    Return:
        If return_lse is True, return a tuple of (output, lse), else return output.
    """

    batch = q.shape[0]
    seq_q = q.shape[1]
    head_num = q.shape[2]
    head_size_v = v_cache_lp.shape[-1]
    tmo_out = torch.empty((q.shape[0], q.shape[1], q.shape[2], head_size_v), dtype=q.dtype, device=q.device) if out is None else out
    out_hp = torch.empty((q.shape[0], q.shape[1], q.shape[2], head_size_v), dtype=q.dtype, device=q.device)
    out_lse = torch.empty((batch, head_num, seq_q), dtype=torch.float32, device=q.device)
    out_lse_hp = torch.empty((batch, head_num, seq_q), dtype=torch.float32, device=q.device)

    torch.ops.torch_mlu_ops.single_query_cached_kv_attn(
            q, k_cache_lp, v_cache_lp, tmo_out, block_tables_lp, context_lens_lp, out_lse,
            k_cache_quant_scale_lp, v_cache_quant_scale_lp, alibi_slopes, max_contxt_len_lp,
            -1, -1, softmax_scale, True, kv_cache_quant_bit_size_lp)
    torch.ops.torch_mlu_ops.single_query_cached_kv_attn(
            q, k_cache_hp, v_cache_hp, out_hp, block_tables_hp, context_lens_hp, out_lse_hp,
            k_cache_quant_scale_hp, v_cache_quant_scale_hp, alibi_slopes, max_contxt_len_hp,
            -1, -1, softmax_scale, True, kv_cache_quant_bit_size_hp)
    torch.ops.torch_mlu_ops.update_out_and_lse(tmo_out, out_lse, out_hp, out_lse_hp, None, None, None)
    return (tmo_out, out_lse) if return_lse else tmo_out

def apply_rotary(
    input: torch.Tensor,
    sin_cache: torch.Tensor,
    cos_cache: torch.Tensor,
    position_ids: Optional[torch.Tensor],
    cu_seqlens: Optional[torch.Tensor],
    interleaved: bool,
    discrete: bool,
    dynamic_ntk: bool,
    max_seqlen: int,
) -> torch.Tensor:
    """
    Apply rotary embedding to the input tensor.

    Math:
        rotary_dim = sin_cache.shape[-1]
        input_rot = input[..., :rotary_dim].clone()
        if interleaved is True:
            x1, x2 = input_rot[..., ::2].clone(), input_rot[..., 1::2].clone()
            input_rot[..., ::2], input_rot[..., 1::2] = -x2, x1
        else:
            x1, x2 = input_rot.chunk(2, dim=-1)
            input_rot = concat((-x2, x1), dim=-1)
        input = input_rot * sin_cache[position_ids] + input * cos_cache[position_ids]
        return input

    Args:
        input (torch.Tensor): Shape is (batch, seq, head_num, head_size) or (total_seq, head_num, head_size).
                              The head_num is usually equal to head_num_q+head_num_k.
        sin_cache (torch.Tensor): If dynamic_ntk is true, the shape must be (batch, rotary_seq_len, rotary_dim),
            otherwise shape must be (rotary_seq_len, rotary_dim).
        cos_cache (torch.Tensor): If dynamic_ntk is true, the shape must be (batch, rotary_seq_len, rotary_dim),
            otherwise shape must be (rotary_seq_len, rotary_dim).
        position_ids (torch.Tensor): The position index of input tokens.
            If discrete is True, shape must be (batch, seq) or (total_seq), which indicates each token has a position index.
            If discrete is False, a tensor shape of (batch) specifies the start position of each batch,
                                  a tensor of None indicates start position is 0 for each batch.
        cu_seqlens (torch.Tensor): The sequence length of each batch in pack mode.
        interleaved (bool): If interleaved is True, apply cross rotary embedding, otherwise apply fold rotary embedding.
        dynamic_ntk (bool): A boolean indicates whether all batches share the same sin_cache and cos_cache.
        max_seqlen (int): The maximum sequence lengths of input.

    Type:
        input: float, half, bfloat16.
        sin_cache: same as input.
        cos_cache: same as input.
        position_ids: int32.
        cu_seqlens: int32.

    Return:
        Return the input tensor after applying rotary embedding.
    """
    torch.ops.torch_mlu_ops.apply_rotary(
        input,
        sin_cache,
        cos_cache,
        position_ids,
        cu_seqlens,
        interleaved,
        discrete,
        dynamic_ntk,
        max_seqlen,
    )
    return input

def reshape_linear_cache(
    key: torch.Tensor,
    value: Optional[torch.Tensor],
    key_cache: torch.Tensor,
    value_cache: Optional[torch.Tensor],
    context_lengths: torch.Tensor,
    max_context_len: int,
    packed: bool,
    context_seq_offset: Optional[torch.Tensor],
    cache_bs_id: Optional[torch.Tensor],
    cache_seqlen_offset: Optional[torch.Tensor],
) -> Union[Tuple[torch.Tensor], torch.Tensor]:
    """
    Put key and value into key_cache and value_cache.

    Math:
        for idx in range(batch):
            key_cache[idx] = key[idx][:, context_length[idx]]
            value_cache[idx] = value[idx][:, context_length[idx]]
    Args:
        key (torch.Tensor): The key tensor. If use pack mode, shape is (batch, seqlen, head_num_kv, head_size),
                            else, shape is (total_seqlen, head_num_kv, head_size).
        value (torch.Tensor): The value tensor. Shape is the same as key.
        key_cache (torch.Tensor): The key_cache tensor. Shape is (max_batch, head_num_kv , cache_mem_len, head_size).
        value_cache (torch.Tensor): The value_cache tensor. Shape is the same as key_cache.
        context_lengths (torch.Tensor): Store key or cache lengths. If packed, shape is (batch + 1), else
                                        shape is (batch).
        max_context_len (int): The maximum sequence length of context.
        packed (bool): A boolean value indicates whether to use pack mode.
        context_seq_offset (torch.Tensor): Store the sequence offset of context. Shape if (batch).
        cache_bs_id (torch.Tensor): Indicate context will be put into which batch of cache. A negative value
                                    indicates that the corresponding batch will be ignored. Shape is (batch).
        cache_seqlen_offset (torch.Tensor): Indicate the offset of the position in the cache. Shape is (batch).
    Type:
        key: float, half, bfloat16.
        value: float, half, bfloat16.
        key_cache: float, half, bfloat16.
        value_cache: float, half, bfloat16.
        context_lengths: int32.
        max_context_len: int32.
        packed: bool.
        context_seq_offset: int32.
        cache_bs_id: int32.
        cache_seqlen_offset: int32.
    Return:
        Support inplace outputs.
        Directly return the given key_cache and value_cache if value_cache existed, else return key_cache only.
    """
    torch.ops.torch_mlu_ops.reshape_linear_cache(
        key,
        value,
        key_cache,
        value_cache,
        context_lengths,
        max_context_len,
        packed,
        context_seq_offset,
        cache_bs_id,
        cache_seqlen_offset,
    )
    return (key_cache, value_cache) if value_cache is not None else key_cache

def reshape_paged_cache(
    k: torch.Tensor,
    v: Optional[torch.Tensor],
    k_cache: torch.Tensor,
    v_cache: Optional[torch.Tensor],
    slot_mapping: torch.Tensor
) -> Union[Tuple[torch.Tensor], torch.Tensor]:
    """
    Perform reshape_paged_cache operation.

    Math:
        for i in range(num_tokens):
            if slot_mapping[i] >= 0:
                block_id = torch.div(slot_mapping[i], block_size, rounding_mode='floor')
                block_offset = slot_mapping[i] % block_size
                k_cache[block_id, :, block_offset, :] = k[i]
                v_cache[block_id, :, block_offset, :] = v[i]

    Args:
        k (torch.Tensor): The key tensor. Shape is (num_token, head_num_kv, head_size).
        v (torch.Tensor): The value tensor. Shape is the same as v.
        k_cache (torch.Tensor): The key_cache tensor. Shape is (block_num, head_num_kv, block_size, head_size).
        v_cache (torch.Tensor): The value_cache tensor. Shape is the same as key_cache.
        slot_mapping (torch.Tensor): The slot_mapping tensor. Shape is (num_token).

    Types:
        k: float, half, bfloat16.
        v: float, half, bfloat16.
        k_cache: float, half, bfloat16.
        v_cache: float, half, bfloat16.
        slot_mapping: int32.

    Return:
        Support inplace outputs.
        Directly return the given k_cache and v_cache if v_cache existed, else return k_cache only.
    """
    torch.ops.torch_mlu_ops.reshape_paged_cache(k, v, k_cache, v_cache, slot_mapping)
    return (k_cache, v_cache) if v_cache is not None else k_cache

def quant_to_paged_cache(
    k: torch.Tensor,
    v: Optional[torch.Tensor],
    k_cache: torch.Tensor,
    v_cache: Optional[torch.Tensor],
    k_cache_quant_scale: torch.Tensor,
    v_cache_quant_scale: Optional[torch.Tensor],
    slot_mapping: torch.Tensor,
) -> Tuple[torch.Tensor]:
    """
    Quant k and v, put quanted_kv into k_cache and v_cache, put quanted_scale into
    k_cache_quant_scale and v_cache_quant_scale.

    Math:
        key_i = key_i[batch_idx]
        max_value = max(key_i.abs(), dim=-1)
        int_max = float(2 ** (quant_bit - 1) - 1)
        scale_i = float(max_value / int_max)
        key_i /= key_i
        k_cache[pos] = k_i
        k_cache_quant_scale[pos] = scale_i

    Args:
        k (torch.Tensor): The key tensor. Shape is [token_nums, head_num_kv, head_size].
        v (torch.Tensor): The value tensor. Shape is the same as k.
        k_cache (torch.Tensor): The k_cache tensor. Shape is [block_nums, head_num_kv, block_size, head_size].
        v_cache (torch.Tensor): The v_cache tensor. Shape is the same as v_cache.
        k_cache_quant_scale (torch.Tensor): The k_cache_quant_scale tensor. Shape is [block_nums, head_num_kv, block_size].
        v_cache_quant_scale (torch.Tensor): The v_cache_quant_scale tensor. Shape is the same as v_cache_quant_scale.
        slot_mapping (torch.Tensor): The slot_mapping tensor. Shape is [token_nums].

    Type:
        k: float, half, bfloat16.
        v: float, half, bfloat16.
        k_cache: int8.
        v_cache: int8.
        k_cache_quant_scale: float.
        v_cache_quant_scale: float.
        slot_mapping: int32.

    Return:
        Support inplace outputs.
        Directly return the given k_cache, v_cache, k_cache_quant_scale and v_cache_quant_scale if v_cache existed,
        else return k_cache and k_cache_quant_scale.
    """
    torch.ops.torch_mlu_ops.quant_to_paged_cache(
        k, v, k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale, slot_mapping
    )
    return (k_cache, v_cache, k_cache_quant_scale, v_cache_quant_scale) if v_cache is not None else (k_cache, k_cache_quant_scale)

def offline_quant_to_paged_cache(
    k: torch.Tensor,
    v: Optional[torch.Tensor],
    k_cache_scale: torch.Tensor,
    v_cache_scale: Optional[torch.Tensor],
    slot_mapping: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: Optional[torch.Tensor],
) -> Union[Tuple[torch.Tensor], torch.Tensor]:
    """
    Apply quant_to_paged_cache operation on key and value.

    Math:
        tokens_num = k.shape[0]
        block_size = k_cache.shape[2]
        for i in range(tokens_num):
            if slot_mapping[i] >= 0:
                key_i = k[i]
                value_i = v[i]
                block_id = torch.div(slot_mapping[i], block_size, rounding_mode='floor')
                block_offset = slot_mapping[i] % block_size
                key_cache_i = k_cache[block_id, :, block_offset, :]
                value_cache_i = v_cache[block_id, :, block_offset, :]
                quant_key_i = quant2int8(key_i, k_cache_scale)
                quant_value_i = quant2int8(value_i, v_cache_scale)
                k_cache_i[...] = quant_key_i
                v_cache_i[...] = quant_value_i

    Args:
        k (torch.Tensor): The key tensor, shape is (num_tokens, num_heads, head_size).
        v (torch.Tensor): The value tensor. shape is (num_tokens, num_heads, head_size).
        k_cache_scale (torch.Tensor): The key_cache_scale tensor. shape is (num_heads, head_size).
        v_cache_scale (torch.Tensor): The value_cache_scale tensor. shape is (num_heads, head_size).
        slot_mapping (torch.Tensor): The slot_mapping tensor. shape is (num_tokens).
        k_cache (torch.Tensor): The key_cache tensor. shape is (num_blocks, num_heads, block_size, head_size).
        v_cache (torch.Tensor): The value_cache tensor. shape is (num_blocks, num_heads, block_size, head_size).

    Type:
        k: float, half, bfloat16.
        v: float, half, bfloat16.
        k_cache_scale: float.
        v_cache_scale: float.
        slot_mapping: int32.
        k_cache: int8.
        v_cache: int8.

    Return:
        Support inplace outputs.
        Directly return the given k_cache and v_cache if v_cache existed, else return k_cache only.

    """
    torch.ops.torch_mlu_ops.offline_quant_to_paged_cache(
        k, v, k_cache_scale, v_cache_scale, slot_mapping, k_cache, v_cache,
    )
    return (k_cache, v_cache) if v_cache is not None else k_cache

def quant_to_linear_cache(
    key: torch.Tensor,
    value: Optional[torch.Tensor],
    key_cache: torch.Tensor,
    value_cache: Optional[torch.Tensor],
    key_cache_quant_scale: torch.Tensor,
    value_cache_quant_scale: Optional[torch.Tensor],
    context_lengths: torch.Tensor,
    max_context_len: int,
    packed: bool,
    context_seq_offset: Optional[torch.Tensor],
    cache_bs_id: Optional[torch.Tensor],
    cache_seqlen_offset: Optional[torch.Tensor],
    quant_bit: int = 8,
) -> Tuple[torch.Tensor]:
    """
    Quant key and value, put quanted_kv into key_cache and value_cache, put quanted_scale into
    key_cache_quant_scale and value_cache_quant_scale.

    Math:
        key_i = key[batch_idx]
        if groupwise quantize:
            key_i = key_i.reshape(head_num, -1, group_size)
        else:
            key_i = key_i.reshape(head_num, -1, head_size)
        max_value = max(key_i.abs(), dim=-1)
        int_max = float(2 ** (quant_bit - 1) - 1)
        scale_i = max_value / int_max
        key_i /= scale_i
        key_i = clip(key_i, -2 ** (quant_bit - 1), 2 ** (quant_bit - 1) - 1)
        if quant_bit == 4:
            key_i = key_i.flatten()
            d0 = key_i[0::2].to(uint8)
            d1 = key_i[0::2].to(uint8)
            dp = (d1 << 4) + (d0 & 0x0f)
            key_i = key_i.to(int8).reshape(head_num, -1, head_size // 2)
        k_cache[pos] = k_i
        k_cache_quant_scale[pos] = scale_i

    Args:
        key (torch.Tensor): The key tensor. If packed, shape is (total_seqlen, head_num_kv, head_size),
                            else, shape is (batch, seqlen, head_num_kv, head_size).
        value (torch.Tensor): The value tensor. Shape is the same as key.
        key_cache (torch.Tensor): The key cache tensor. Shape is (max_batch, head_num_kv, cache_mem_len, head_size).
        value_cache (torch.Tensor): The value cache tenosr. Shape is the same as key_value.
        key_cache_quant_scale (torch.Tensor): The key_cache_quant_scale tensor. If use groupwise quantization, shape is
                                              (max_batch, head_num_kv, cache_mem_len, group_num), otherwise shape is
                                              (max_batch, head_num_kv, cache_mem_len).
        value_cache_quant_scale (torch.Tensor): The value_cache_quant_scale tensor. Shape is the same as key_cache_quant_scale.
        context_lengths (torch.Tensor): The context_lengths tensor. If packed, shape is (batch + 1), else, shape is (batch).
        max_context_len: The maximum sequence length of context.
        packed: A boolean value indicates whether to use pack mode.
        context_seq_offset (torch.Tensor): Indicate the sequence offset of context, shape is (batch).
        cache_bs_id (torch.Tensor): Indicate key and value will be put into which batch of kv_cache. A negative value indicates
                                    that the corresponding batch will be ignored. Shape is (batch).
        cache_seqlen_offset (torch.Tensor): Indicate the offset of the position in the cache. Shape is (batch).

    Types:
        key: float, half, bfloat16.
        value: float, half, bfloat16.
        key_cache: int8.
        value_cache: int8.
        key_cache_quant_scale: float32.
        value_cache_quant_scale: float32.
        context_lengths: int32.
        max_context_len: int32.
        packed: bool.
        context_seq_offset: int32.
        cache_bs_id: int32.
        cache_seqlen_offset: int32.

    Return:
        Support inplace outputs.
        Directly return the given key_cache, value_cache, key_cache_quant_scale and value_cache_quant_scale if value_cache existed,
        else return key_cache and key_cache_quant_scale.
    """
    torch.ops.torch_mlu_ops.quant_to_linear_cache(
        key,
        value,
        key_cache,
        value_cache,
        key_cache_quant_scale,
        value_cache_quant_scale,
        context_lengths,
        max_context_len,
        packed,
        context_seq_offset,
        cache_bs_id,
        cache_seqlen_offset,
        quant_bit,
    )
    return (key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale) if value_cache is not None else (key_cache, key_cache_quant_scale)

def offline_quant_to_linear_cache(
    key: torch.Tensor,
    value: Optional[torch.Tensor],
    key_cache: torch.Tensor,
    value_cache: Optional[torch.Tensor],
    key_cache_quant_scale: torch.Tensor,
    value_cache_quant_scale: Optional[torch.Tensor],
    context_lengths: torch.Tensor,
    max_context_len: int,
    quant_mode: int,
    packed: bool,
    context_seq_offset: Optional[torch.Tensor],
    cache_bs_id: Optional[torch.Tensor],
    cache_seqlen_offset: Optional[torch.Tensor],
) -> Tuple[torch.Tensor]:
    """
    Quantize current key and value, Then store key and value to key_cache and value_cache.

    Math:
        key_i = key[batch_idx].transpose(1, 0)
        if quant_mode == 0:
            key_i /= key_cache_quant_scale.reshape(head_num, 1, head_size)
        else:
            key_i /= key_cache_quant_scale.reshape(head_num, seq, 1)
        key_i = clip(round(key_i), -128, 127)
        key_cache[pos] = key_i

    Args:
        key (torch.Tensor): The key tensor. If packed, shape is (total_seqlen, head_num_kv, head_size),
                            else, shape is (batch, seqlen, head_num_kv, head_size).
        value (torch.Tensor): The value tensor. Shape is the same as key.
        key_cache (torch.Tensor): The key_cache tensor. Shape is (max_batch, head_num, cache_mem_len, head_size).
        value_cache (torch.Tensor): The value_cache tensor. Shape is the same as key_cache.
        key_cache_quant_scale (torch.Tensor): If per_channel quantize, shape is (head_num_kv, head_size).
                                              If per_head quantize, shape is (head_num_kv, cache_mem_len).
        value_cache_quant_scale (torch.Tensor): Shape is the same as key_cache_quant_scale.
        context_lengths (torch.Tensor): A tensor indicate context lengths. Shape is (batch+1) if packed is True, which
                                        stores the cumsum of seqlen, otherwise shape is (batch).
        max_context_len: The maximum sequence length of context.
        quant_mode: A int value indicates which quantize mode will be used. If quant_mode is 0, use per_channel
                    quantize, otherwise use per_head quantize.
        packed: A boolean value indicates whether to use pack mode.
        context_seq_offset (torch.Tensor): Indicate the sequence offset of context, shape is (batch).
        cache_bs_id (torch.Tensor): Indicate context will be put into which batch of cache. A negative value
                                    indicates that the corresponding batch will be ignored. Shape is (batch).
        cache_seqlen_offset (torch.Tensor): Indicate the offset of the position in the cache. Shape is (batch).

    Types:
        key: float, half, bfloat16.
        value: float, half, bfloat16.
        key_cache: int8.
        value_cache: int8.
        key_cache_quant_scale: float32.
        value_cache_quant_scale: float32.
        context_lengths: int32.
        max_context_len: int32.
        quant_mode: int32.
        packed: bool.
        context_seq_offset: int32.
        cache_bs_id: int32.
        cache_seqlen_offset: int32.

    Return:
        Support inplace outputs.
        Directly return the given key_cache, value_cache, key_cache_quant_scale and value_cache_quant_scale if value_cache existed,
        else return key_cache and key_cache_quant_scale.
    """
    torch.ops.torch_mlu_ops.offline_quant_to_linear_cache(
        key,
        value,
        key_cache,
        value_cache,
        key_cache_quant_scale,
        value_cache_quant_scale,
        context_lengths,
        max_context_len,
        quant_mode,
        packed,
        context_seq_offset,
        cache_bs_id,
        cache_seqlen_offset,
    )
    return (key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale) if value_cache is not None else (key_cache, key_cache_quant_scale)

def swap_blocks(dst: torch.Tensor,
                src: torch.Tensor,
                block_mapping: Dict[int, int]) -> torch.Tensor:
    """
    Copy src value to dst according to block_mapping.

    Math:
        for key, value in block_mapping.items():
            dst[value] = src[key]

    Args:
        dst (torch.Tensor): Destination tensor. Shape is (num_blocks, num_heads, block_size, head_size).
        src (torch.Tensor): Source tensor. Shape is (num_blocks, num_heads, block_size, head_size).
        block_mapping (dict): Mapping table of src and dst.

    Types:
        dst: Unlimited.
        src: Unlimited.
        block_mapping: [int, int].

    Return:
        Support inplace outputs.
        Directly return the given dst.
    """
    torch.ops.torch_mlu_ops.swap_blocks(dst, src, block_mapping)
    return dst

def copy_blocks(
    k_caches: List[torch.Tensor],
    v_caches: Optional[List[torch.Tensor]],
    block_mapping: Dict[int, List[int]]
) -> Union[Tuple[List[torch.Tensor], List[torch.Tensor]], List[torch.Tensor]]:
    """
    Remap the k_caches and v_caches with the block indexes pairs specified by block_mapping.

    Math:
        for src, dsts in block_mapping.items():
            srcs = [src for i in range(len(dsts))]
            for key_cache in k_caches:
                key_cache[dsts] = key_cache[srcs]
            for value_cache in v_caches:
                value_cache[dsts] = value_cache[srcs]

    Args:
        k_caches (List[torch.Tensor]): A tensor list of k_cache.
        v_caches (List[torch.Tensor]): A tensor list of v_cache.
        block_mapping (Dict[int, List[int]]): The pairs of source and destination block indexes.

    Type:
        k_caches: int8, uint8, int16, int32, int64, half, float, bfloat16.
        v_caches: int8, uint8, int16, int32, int64, half, float, bfloat16.

    Return:
        Support inplace outputs.
        Directly return the given k_caches and v_caches.
    """
    torch.ops.torch_mlu_ops.copy_blocks(k_caches, v_caches if v_caches is not None else [], block_mapping)
    return (k_caches, v_caches) if v_caches is not None else k_caches

def quant_matmul(
    a_tensor: torch.Tensor,
    a_scale: Optional[torch.Tensor],
    a_zero: Optional[torch.Tensor],
    b_tensor: torch.Tensor,
    b_scale: Optional[torch.Tensor],
    b_zero: Optional[torch.Tensor],
    bias: Optional[torch.Tensor],
    c_tensor: Optional[torch.Tensor],
    c_scale: Optional[torch.Tensor],
    c_zero: Optional[torch.Tensor],
    gemm_output_scale: Optional[torch.Tensor],
    gemm_output_zero: Optional[torch.Tensor],
    data_type: Optional[str],
    quant_algo: str,
    a_quant_layout: str,
    b_quant_layout: str,
    quant_bit_size: int = 8,
    act_mode: str = "none",
    use_hp_active: bool = False,
    act_coef: float = 1.0,
    alpha: float = 1.0,
    beta: float = 1.0,
    trans_a: bool = False,
    trans_b: bool = True,
) -> torch.Tensor:
    """
    Full-featured function to do quantized matmul.
    It is recommended to use the function with fewer parameters such as 'weight_only_quant_matmul' or 'smooth_quant_matmul'.
    """
    return torch.ops.torch_mlu_ops.quant_matmul(
        a_tensor,
        a_scale,
        a_zero,
        b_tensor,
        b_scale,
        b_zero,
        bias,
        c_tensor,
        c_scale,
        c_zero,
        gemm_output_scale,
        gemm_output_zero,
        data_type,
        None,
        quant_algo,
        a_quant_layout,
        b_quant_layout,
        quant_bit_size,
        act_mode,
        use_hp_active,
        act_coef,
        alpha,
        beta,
        trans_a,
        trans_b,
    )

def quant_matmul_allreduce(
    cncl_comm,
    a_tensor: torch.Tensor,
    a_scale: Optional[torch.Tensor],
    a_zero: Optional[torch.Tensor],
    b_tensor: torch.Tensor,
    b_scale: Optional[torch.Tensor],
    b_zero: Optional[torch.Tensor],
    bias: Optional[torch.Tensor],
    c_tensor: Optional[torch.Tensor],
    c_scale: Optional[torch.Tensor],
    c_zero: Optional[torch.Tensor],
    gemm_output_scale: Optional[torch.Tensor],
    gemm_output_zero: Optional[torch.Tensor],
    data_type: Optional[str],
    quant_algo: str,
    a_quant_layout: str,
    b_quant_layout: str,
    quant_bit_size: int = 8,
    alpha: float = 1.0,
    beta: float = 1.0,
    trans_a: bool = False,
    trans_b: bool = True,
    block_m: int = 0
) -> torch.Tensor:
    return torch.ops.torch_mlu_ops.quant_matmul_allreduce(
        cncl_comm,
        a_tensor,
        a_scale,
        a_zero,
        b_tensor,
        b_scale,
        b_zero,
        bias,
        c_tensor,
        c_scale,
        c_zero,
        gemm_output_scale,
        gemm_output_zero,
        data_type,
        None,
        quant_algo,
        a_quant_layout,
        b_quant_layout,
        quant_bit_size,
        alpha,
        beta,
        trans_a,
        trans_b,
        block_m
    )

def active(input: torch.Tensor, act_mode: str, is_gated: bool, active_coef: float = 1.0) -> torch.Tensor:
    """
    Apply activation to the input tensor.

    Math:
        C = input.shape[-1]
        if is_gated = True:
            output = active(input[..., :C//2]) * input[..., C//2:]
        else:
            output = active(input)

    Args:
        input (torch.Tensor): Shape is (..., C)
        act_mode (str): The activation mode, must be 'silu', 'gelu', 'swish' or 'quick_gelu'.
        is_gated (bool): If use gated activation.
        active_coef (float): The coefficient used in the swish activation. Default is 1.0.

    Type:
        input: float, half, bfloat16.
        output: same as input.

    Return:
        A tensor with the same shape and dtype of input.
    """
    out_shape = list(input.shape)
    if is_gated:
        out_shape[-1] = input.size(-1) // 2
    output = torch.empty(out_shape, dtype=input.dtype, device=input.device)
    torch.ops.torch_mlu_ops.active(input, output, None, None, act_mode,
                                          is_gated, 0, 0, active_coef)
    return output

def fused_moe(
        hidden_states: torch.Tensor,
        gating_output: torch.Tensor,
        w1: torch.Tensor,
        w2: torch.Tensor,
        bias1: Optional[torch.Tensor],
        bias2: Optional[torch.Tensor],
        residual: Optional[torch.Tensor],
        input_smooth: Optional[torch.Tensor],
        act_smooth: Optional[torch.Tensor],
        w1_scale: Optional[torch.Tensor],
        w2_scale: Optional[torch.Tensor],
        topk: int,
        renormalize: bool,
        gated: bool,
        act_mode: str,
        start_expert_id: int = 0,
        block_n: int = 0,
        cncl_comm: int = 0,
        w1_quant_flag: Optional[List[int]] = None,
        w2_quant_flag: Optional[List[int]] = None,
    ) -> torch.Tensor:
    """
    Apply mixtrue of experts operation to the input tensor.

    Math:
        Please refer to 'test/pytest/test_moe.py' in this repository.

    Args:
        hidden_states (torch.Tensor): Shape is (batch, seq, hidden_size) or (total_seq, hidden_size).
        gating_output (torch.Tensor): The output of router FC. Shape is (batch, seq, num_expert) or (total_seq, num_expert).
        w1 (torch.Tensor): shape are shown as followed:
            (expert_size, inner_size * (1 + gated), hidden_size) for normal mode.
            (expert_size, inner_size * (1 + gated), hidden_size / 2) for int4 quantization.
        w2 (torch.Tensor): shape are shown as followed:
            (expert_size, hidden_size, inner_size) for normal mode.
            (part_a, hidden_size, part_b, inner_size) for retiled mode, which part_a * part_b must be expert_size.
            (expert_size, hidden_size, inner_size / 2) for int4 quantization.
        bias1 (torch.Tensor): Must be None.
        bias2 (torch.Tensor): Must be None.
        residual (torch.Tensor): Shape is (batch, seq, hidden_size) or (total_seq, hidden_size).
        input_smooth (torch.Tensor): The smooth scale of hidden_states, shape is (expert_size, hidden_size).
        act_smooth (torch.Tensor): The smooth scale of the activation output, shape is (expert_size, inner_size).
        w1_scale (torch.Tensor): The per-channel scale of w1, shape are shown as followed:
            (expert_size, (1 + gated) * inner_size) for normal quantization.
            (w1_quant_group, expert_size, (1 + gated) * inner_size) for groupwise quantization.
        w2_scale (torch.Tensor): The per-channel scale of w2, shape are shown as followed:
            (expert_size, hidden_size) for normal quantization.
            (w2_quant_group, expert_size, hidden_size) for groupwise quantization.
        topk (int): The number of experts selected from the total experts, topk <= num_expert.
        renormalize (bool): If do normalization on the topk_logits.
        gated (bool): If use gated activation.
        act_mode (str): The activation mode, must be 'silu' or 'gelu'.
        start_expert_id (int): The starting id of the expert.
        block_n (int): Take effect only if block_n > 0, otherwise the recommended value will be used.
        cncl_comm (int): The handle of CnclComm.
        w1_quant_flag (List[int]): A list of int values, the elements must be 4 or 8, which specify the quantization bits of each group,
            the other values may cause undefined behavior. Length is (expert_size * w1_quant_group).
        w2_quant_flag (List[int]): A list of int values, the elements must be 4 or 8, which specify the quantization bits of each group,
            the other values may cause undefined behavior. Length is (expert_size * w2_quant_group).

    Type:
        hidden_states: float, half, bfloat16.
        gating_output: same as hidden_states.
        w1: float, half, bfloat16, int8, int8(int4x2).
        w2: same as w1.
        bias1: same as hidden_states.
        bias2: same as hidden_states.
        residual: same as hidden_states.
        input_smooth: float.
        act_smooth: float.
        w1_scale: float.
        w2_scale: float.
        output: same as hidden_states.

    Return:
        A tensor with the same shape of hidden_states.

    Note:
        The elements in w1_quant_flag and w2_quant_flag must be 4 or 8.
    """
    return torch.ops.torch_mlu_ops.fused_moe(
        hidden_states,
        gating_output,
        w1,
        w2,
        bias1,
        bias2,
        residual,
        input_smooth,
        act_smooth,
        w1_scale,
        w2_scale,
        w1_quant_flag,
        w2_quant_flag,
        topk,
        renormalize,
        gated,
        act_mode,
        start_expert_id,
        block_n,
        cncl_comm
    )

def matmul(
        a: torch.Tensor,
        b: torch.Tensor,
        bias: Optional[torch.Tensor] = None,
        c: Optional[torch.Tensor] = None,
        act_mode: str = 'none',
        alpha: float = 1.0,
        beta: float = .0,
        fast_act: bool = True,
        approximate: bool = True,
        dtype: torch.dtype = None,
        a_scale: float = 1.0,
        b_scale: float = 1.0,
        trans_a: bool = False,
        trans_b: bool = True,
    ) -> torch.Tensor:
    """
    Perform matrix multiplication operation.

    Math:
        ab = matmul(a, b)
        if c is not None:
            ab = alpha * ab + beta * c
        if bias is not None:
            ab += bias
        if act_mode is not 'none':
            ab = active(ab)
        output = ab

    Args:
        a (torch.Tensor): If trans_a, shape is (k, m), else shape is (m, k).
        b (torch.Tensor): If trans_b, shape is (n, k), else shape is (k, n).
        bias (torch.Tensor): Shape is (n).
        c (torch.Tensor): Shape is (m, n).
        act_mode (str): The activation type, must be 'gelu', 'relu', 'silu' or 'none'.
        alpha (float): The coefficient of ab.
        beta (float): The coefficient of c.
        fast_act (bool): Describing the algorithm that used in the implementation of the activation function.
             When the value is true, use the fastest algorithm of activation, otherwise use the high-precision algorithm.
        approximate(bool): Describes the implementation logic of different GELU approximation algorithms.
        d (torch.Tensor): Shape is (m, n).
        a_scale (float): quant_scale of a.
        b_scale (float): quant_scale of b.

    Type:
        a: float, half, bfloat16, int8.
        b: same as a.
        bias: float, half, bfloat16.
        c: float, half, bfloat16.
        output: float, half, bfloat16.

    Return:
        Return output tensor.
    """
    d_dtype = None if dtype is None else _torchDtype2Str(dtype)
    return torch.ops.torch_mlu_ops.matmul(a, b, None, bias, c, d_dtype, act_mode, alpha, beta,
                                           fast_act, approximate, a_scale, b_scale, trans_a, trans_b)

def batch_matmul(
        a: torch.Tensor,
        b: torch.Tensor,
        c: Optional[torch.Tensor] = None,
        alpha: float = 1.0,
        beta: float = .0,
        a_scale: float = 1.0,
        b_scale: float = 1.0,
        trans_a: bool = False,
        trans_b: bool = True,
    ) -> torch.Tensor:
    """
    Perform matrix multiplication operation.

    Math:
        ab = batch_matmul(a, b)
        if c is not None:
            c = alpha * ab + beta * c
            output = c
        else:
            output = ab

    Args:
        a (torch.Tensor): Shape is (batch, m, k).
        b (torch.Tensor): Shape is (batch, n, k).
        c (torch.Tensor): Shape is (batch, m, n), c will be changed when the calculation is completed.
        alpha (float): The coefficient of ab.
        beta (float): The coefficient of c.
        a_scale (float): quant_scale of a.
        b_scale (float): quant_scale of b.

    Type:
        a: float, half, bfloat16, int8.
        b: same as a.
        c: float, half, bfloat16.
        output: float, half, bfloat16.

    Return:
        Return output tensor.
    """
    tmo_out = c
    if c is None:
        assert a.dtype in (torch.half, torch.float, torch.bfloat16), "dtype must be half, float or bfloat16, when c is None."
        m = a.shape[-1] if trans_a else a.shape[-2]
        n = b.shape[-2] if trans_b else b.shape[-1]
        batch = a.shape[0]
        tmo_out = torch.zeros((batch, m, n), dtype=a.dtype, device=a.device)
    torch.ops.torch_mlu_ops.batch_matmul(a, b, tmo_out, alpha, beta, a_scale, b_scale, trans_a, trans_b)
    return tmo_out

def matmul_allreduce(cncl_comm,
                     a: torch.Tensor,
                     b: torch.Tensor,
                     bias: Optional[torch.Tensor] = None,
                     c: Optional[torch.Tensor] = None,
                     alpha: float = 1.0,
                     beta: float = .0,
                     block_m: int = 0
                    ):
    """
    Parallelly apply matrix multiplication with TP-allreduce.

    Math:
        ab = matmul(a, b)
        if c is not None:
            ab = alpha * ab + beta * c
        if bias is not None:
            ab += bias
        output = ab
        output_allreduce = allreduce(output)

    Args:
        cncl_comm (int): The handle of CnclComm.
        a (torch.Tensor): Shape is (m, k).
        b (torch.Tensor): Shape is (n, k).
        bias (torch.Tensor): Shape is (n).
        c (torch.Tensor): Shape is (m, n).
        alpha (float): The coefficient of ab.
        beta (float): The coefficient of c.
        block_m (int): If the value is greater than 0, the specified value will be used; if the value is less than 0, the recommended value will be used.

    Type:
        a: float, half, bfloat16.
        b: same as a.
        bias: same as a.
        c: same as a.
        output: same as a.

    Return:
        Return output tensor.
    """
    return torch.ops.torch_mlu_ops.matmul_allreduce(cncl_comm, a, b, bias, c, None, alpha, beta, block_m)

def group_gemm(a: torch.Tensor,
               b: torch.Tensor,
               m_list: torch.Tensor,
               expand_idx: Optional[torch.Tensor],
               c: Optional[torch.Tensor],
               alpha: Optional[torch.Tensor],
               beta: Optional[torch.Tensor],
               max_m: int,
               bias: Optional[torch.Tensor] = None,
               ) -> torch.Tensor:
    '''
    Perform the grouped matrix multiplication operation.

    Math:
        if expand_idx:
            a = a[expand_idx]
        a_list = a.split(m_list)
        c_list = c.split(m_list)
        d_list = d.split(m_list)
        experts = len(m_list)
        for i in range(experts):
            d_list[i] = matmul(a_list[i], b[i]) * alpha[i] + beta[i] * c_list[i] + bias[i]
        d = concat(d_list, dim=0)

    Args:
        a (torch.Tensor): Shape is (m, k). If expand_idx exists, total_m = m_list.sum(), otherwise total_m = m.
        b (torch.Tensor): Shape is (experts, n, k).
        m_list (torch.Tensor): Shape is (experts).
        expand_idx (torch.Tensor): optional. Shape is (total_m).
        c (torch.Tensor): optional. Shape is (total_m, n).
        alpha (torch.Tensor): optional. Shape is (experts).
        beta (torch.Tensor): optional. Shape is (experts).
        max_m (int): Maximum possible value in m_list.
        bias (torch.Tensor): optional. Shape is (experts, n).

    Type:
        a: half, bfloat16.
        b: same as a.
        m_list: int32.
        expand_idx: int32.
        c: same as a.
        alpha: float.
        beta: float.
        d: same as a.
        bias: same as a.

    Return:
        Return d.
    '''
    return torch.ops.torch_mlu_ops.group_gemm(a, b, m_list, expand_idx, c, alpha, beta, None, None, bias, _torchDtype2Str(a.dtype), None, None, max_m)

def smooth_quant_group_gemm(a: torch.Tensor,
                            b: torch.Tensor,
                            m_list: torch.Tensor,
                            expand_idx: Optional[torch.Tensor],
                            c: Optional[torch.Tensor],
                            alpha: Optional[torch.Tensor],
                            beta: Optional[torch.Tensor],
                            a_scale: torch.Tensor,
                            b_scale: torch.Tensor,
                            dtype,
                            max_m: int,
                            bias: Optional[torch.Tensor] = None,
                            quant_flag: Optional[List[int]] = None
                            ) -> torch.Tensor:
    '''
    Perform smooth quantized group_gemm operation.

    Math:
        if expand_idx:
            a = a[expand_idx]
        a_list = a.split(m_list)
        c_list = c.split(m_list)
        d_list = d.split(m_list)
        a_scale_list = a_scale.split(m_list)
        experts = len(m_list)
        for i in range(experts):
            ab_scale = a_scale_list[i].outer(b_scale[i])
            ab = matmul(a_list[i], b[i]) * ab_scale
            d_list[i] = ab * alpha[i] + beta[i] * c_list[i] + bias[i]
        d = concat(d_list, dim=0)

    Args:
        a (torch.Tensor): Shape is (m, k). If expand_idx exists, total_m = m_list.sum(), otherwise total_m = m.
        b (torch.Tensor): Shape is (experts, n, k), or (experts, n, k // 2) for int 4 quantization.
        m_list (torch.Tensor): Shape is (experts).
        expand_idx (torch.Tensor): optional. Shape is (total_m).
        c (torch.Tensor): optional. Shape is (total_m, n).
        alpha (torch.Tensor): optional. Shape is (experts).
        beta (torch.Tensor): optional. Shape is (experts), the value can only be 0 or 1.
        a_scale (torch.Tensor): Shape is (total_m).
        b_scale (torch.Tensor): Shape is (experts, n), or (quant_group, experts, n) for groupwise quantization.
        max_m (int): Maximum possible value in m_list.
        dtype (torch.dtype). Specify the data type of output, must be torch.half or torch.bfloat16.
        bias (torch.Tensor): optional. Shape is (experts, n).
        quant_flag (List[int]): A list of int values, the elements must be 4 or 8, which specify the quantization bits of each group,
            the other values may cause undefined behavior. Length is (experts * quant_group).

    Type:
        a: int8.
        b: int8, int8(int4x2).
        m_list: int32.
        expand_idx: int32.
        c: half, bfloat16.
        alpha: float.
        beta: float.
        a_scale: float.
        b_scale: float.
        d: half, bfloat16.
        bias: half, bfloat16.

    Return:
        Return d.

    Note:
        The elements in quant_flag must be 4 or 8.
    '''
    return torch.ops.torch_mlu_ops.group_gemm(a, b, m_list, expand_idx, c, alpha, beta, a_scale, b_scale, bias, _torchDtype2Str(dtype), quant_flag, None, max_m)

def preload(
    weight: torch.Tensor,
    size: int,
) -> None:
    torch.ops.torch_mlu_ops.preload(weight, size)
    return weight

def flash_attn_sq_mm_allreduce(cncl_comm: int,
                               q: torch.Tensor,
                               k: torch.Tensor,
                               v: torch.Tensor,
                               cu_seq_lens_q: Optional[torch.Tensor],
                               cu_seq_lens_kv: Optional[torch.Tensor],
                               alibi_slope: Optional[torch.Tensor],
                               attn_bias: Optional[torch.Tensor],
                               smooth: torch.Tensor,
                               weight: torch.Tensor,
                               weight_scale: torch.Tensor,
                               bias: Optional[torch.Tensor],
                               max_seq_len_q: int,
                               max_seq_len_kv: int,
                               softmax_scale: float,
                               is_causal: bool,
                               window_size_left: int = -1,
                               window_size_right: int = -1,
                               compute_dtype: torch.dtype = torch.float,
                               block_seq: int = 0) -> torch.Tensor:
    """
    Parallelly apply attention-pertoken_smooth_quant-smooth_quant_matmul with TP-allreduce.

    Math:
        attn_out = flash_attn(q, k, v, cu_seq_lens_q, cu_seq_lens_kv)
        quant_a, scale_a = per_token_smooth_quantize(attn_out, smooth)
        output = smooth_quant_matmul(quant_a, scale_a, weight, weight_scale, bias)
        output_allreduce = allreduce(output)

    Args:
        cncl_comm (int): The handle of CnclComm.
        q (torch.Tensor): The query tensor. Shape is (batch, seq_q, head_num_q, head_size) or (total_seq_q, head_num_q, head_size).
        k (torch.Tensor): The key tensor. Shape is (batch, seq_kv, head_num_kv, head_size) or (total_seq_kv, head_num_kv, head_size).
        v (torch.Tensor): The value tensor. Shape is (batch, seq_kv, head_num_kv, head_size) or (total_seq_kv, head_num_kv, head_size).
        cu_seq_lens_q (torch.Tensor): The cusum of seq_q if q is 3-D tensor. Shape is (batch+1).
        cu_seq_lens_kv (torch.Tensor): The cusum of seq_kv if k/v is 3-D tensor. Shape is (batch+1).
        alibi_slope (torch.Tensor): Reserved parameter, must be None.
        attn_bias (torch.Tensor): Reserved parameter, must be None.
        smooth (torch.Tensor): The smooth scale. Shape is (head_num_q*head_size).
        weight (torch.Tensor): The weight of smooth_quant_matmul. Shape is (head_num_q*head_size*TP_NUM, head_num_q*head_size).
        bias (torch.Tensor): The bias of smooth_quant_matmul. Shape is (head_num_q*head_size*TP_NUM).
        max_seq_len_q (int): The maximum value of seq_q.
        max_seq_len_kv (int): The maximum value of seq_kv.
        softmax_scale (int): The scale factor multiplied to qk.
        is_causal (bool): Controlling if use causal mask.
        window_size_left (int): Reserved parameter, must be -1.
        window_size_right (int): Reserved parameter, must be -1.
        compute_dtype (torch.dtype): The dtype used to calculate softmax.
        block_seq (int): Reserved parameter.

    Type:
        q: half, bfloat16.
        k: same as q.
        v: same as q.
        cu_seq_lens_q: int32.
        cu_seq_lens_kv: int32.
        smooth: float.
        weight: int8.
        weight_scale: float.
        bias: same as q.

    Return:
        Return the result of allreduce.

    Note:
        batch must be 1.
    """
    return torch.ops.torch_mlu_ops.flash_attn_sq_mm_allreduce(cncl_comm, q, k, v,
                cu_seq_lens_q, cu_seq_lens_kv, alibi_slope, attn_bias,
                smooth, weight, weight_scale, bias, max_seq_len_q, max_seq_len_kv,
                softmax_scale, is_causal, window_size_left, window_size_right,
                _torchDtype2Str(compute_dtype), block_seq)

def moe_cast_gating(input: torch.Tensor,
                    weight: torch.Tensor) -> None:
    """
    Cast input data type to torch.float32, and perform gating operation.

    Math:
        input_fp32 = input.to(torch.float32).
        gating_output = F.Linear(input_fp32, weight).

    Args:
        input (torch.Tensor): The input tensor, shape is (..., hidden_size), hidden_size must be less than or equal to 16384.
        weight (torch.Tensor): The input tensor, shape is (expert_num, hidden_size), expert_num must be less than or equal to 128.

    Types:
        input: half, bfloat16.
        weight: float.

    Return:
        return a float tensor with shape [..., expert_num].
    """

    return torch.ops.torch_mlu_ops.moe_cast_gating(input, weight)

def moe_softmax_topk(input: torch.Tensor,
                     topk: int,
                     normalize: bool = False,
                     num_expert_group: int = -1,
                     topk_group: int = 0,
                     mask: Optional[torch.Tensor] = None,
                     normed_by : str = "topk_logit") -> Tuple[torch.Tensor]:
    """
    Apply grouped softmax-topk to input tensor.

    Math:
        if num_expert_group <= 1:
            if origin_mask is not None:
                softmax = softmax * origin_mask
            reduce_weight, expert_id = torch.topk(softmax, k=topk, dim=-1)
            if normalize:
                if normed_by == "topk_logit":
                    reduce_weight = reduce_weight / reduce_weight.sum(dim=-1, keepdim=True)
                if normed_by == "softmax_logit":
                    reduce_weight = reduce_weight / softmax.sum(dim=-1, keepdim=True)
            return reduce_weight, expert_id
        else:
            group_size = softmax.shape[-1] // num_expert_group
            new_shape = softmax.shape[:-1] + (num_expert_group, group_size)
            group_data = softmax.view(new_shape)
            group_max_value = group_data.max(dim=-1).values
            group_idx = torch.topk(group_max_value, k=topk_group, dim=-1)[1]
            mask_shape = softmax.shape[:-1] + (num_expert_group,)
            mask = torch.zeros((mask_shape), dtype = torch.bool, device = group_idx.device)
            mask.scatter_(-1, group_idx, True)
            mask = mask.unsqueeze(-1).expand(new_shape)
            masked_data = group_data.masked_fill(~mask, 0.0)
            masked_data = masked_data.reshape(softmax.shape)
            reduce_weight, expert_id = torch.topk(masked_data, k=topk, dim=-1)
            if normalize:
                if normed_by == "topk_logit":
                    reduce_weight = reduce_weight / reduce_weight.sum(dim=-1, keepdim=True)
                if normed_by == "softmax_logit":
                    reduce_weight = reduce_weight / softmax.sum(dim=-1, keepdim=True)
            return reduce_weight, expert_id

    Args:
        input (torch.Tensor): The input tensor, shape is (..., expert_num).
        topk (int): The number of experts that each token would be dispatched to.
        normalize (bool): If do normalization to the reduce_weight after topk operation.
        num_expert_group(int): The number of groups that each expert would be grouped to.
        topk_group(int): The number of groups to be selected from num_expert_group.
        mask (torch.Tensor): The mask multiplied to the softmax result. The dimension must be
        the same as input. The last two dims must be the same as input and the other dims must
        be 1.
        normed_by (str): The mode of normalization, which can be "topk_logit" or "softmax_logit"

    Type:
        input: float, half, bfloat16.
        mask: same as input

    Return:
        Return the reduce_weight and expert_id tensor.

    Note:
        Input and mask must be contiguous.
    """
    outputs = torch.ops.torch_mlu_ops.moe_softmax_topk(input,
                                                       topk,
                                                       num_expert_group,
                                                       topk_group,
                                                       normalize,
                                                       mask,
                                                       normed_by)
    return tuple(outputs)

def moe_expand_input(input: torch.Tensor,
                     gather_idx: torch.Tensor,
                     cusum_token_count: Optional[torch.Tensor] = None,
                     start_expert_id: int = 0,
                     expert_size: int = 0) -> torch.Tensor:
    """
    Expands the input tensor with the indices specified by gather_idx.

    Math:
        if cusum_token_count is None:
            output = input[gather_idx]
        else:
            idx = gather_idx[cusum_token_count[start_expert_id]:cusum_token_count[start_expert_id+expert_count]]
            output = input[idx]

    Args:
        input (torch.Tensor): The input tensor, shape is (token_num, hidden_size).
        gather_idx (torch.Tensor): The input tensor containing the indices to index. Shape is (expand_token_num).
        cusum_token_count (torch.Tensor): The input tensor storing the prefix sum of token count. Shape is (expert_num + 1).
        output (torch.Tensor): The output tensor storing expanded input. Shape is (expand_token_num, hidden_size).
        start_expert_id: the first expert id.
        expert_size: the number of experts currently being processed.

    Types:
        input: int8, float, half, bfloat16.
        gather_idx: int32.
        cusum_token_count: int32.
        output: float, half, bfloat16.
        start_expert_id: int32.
        expert_size: int32.

    Return:
        return the input selected with gather idx.
    """
    return torch.ops.torch_mlu_ops.moe_expand_input(input, gather_idx, cusum_token_count,
                                                    start_expert_id, expert_size)

def moe_gen_idx(expert_id: torch.Tensor,
                expert_num: int) -> Tuple[torch.Tensor]:
    """
    Generate expand_idx, combine_idx, token_count, and cusum_token_count.

    Args:
        expert_id (torch.Tensor): The input tensor stores the expert id of each token, the shape must be [num_token, topk].
        expert_num: the number of expert.

    Types:
        expert_id: int32.
        expert_num: int32.

    Return:
        return a tensor list includes expand_idx, combine_idx, token_count and cusum_token_count.
    """

    outputs = torch.ops.torch_mlu_ops.moe_gen_idx(expert_id, expert_num)
    return tuple(outputs)

def moe_combine_result(input: torch.Tensor,
                       reduce_weight: torch.Tensor,
                       gather_ids: torch.Tensor,
                       residual: Optional[torch.Tensor],
                       cusum_token_count: Optional[torch.Tensor],
                       start_expert_id: int,
                       expert_size: int,
                       bias: Optional[torch.Tensor] = None) -> torch.Tensor:
    '''
    Perform combine result operation in moe.

    Math:
        if cusum_token_count:
            input = input[gather_ids - cusum_token_count[start_expert_id + expert_size]]
        else:
            input = input[gather_ids]

        if bias:
            for i in range(start_expert_id : start_expert_id + expert_size):
                out[cusum_token_count[i] : cusum_token_count[i+1]] += bias[i]

        if cusum_token_count:
            reduce_weight *= (gather_ids >= cusum_token_count[start_expert_id]) *
                             (gather_ids < cusum_token_count[start_expert_id + expert_size])
        out = input * reduce_weight
        out = out.sum(1)
        if residual:
          out += residual

    Args:
        input (torch.Tensor): Shape is (num_tokens, hidden_size).
        reduce_weight (torch.Tensor): Shape is (num_token, topk).
        gather_ids (torch.Tensor): Shape is (num_tokens).
        residual (torch.Tensor): optional. Shape is (num_token, hidden_size).
        bias (torch.Tensor): optional. Shape is (num_expert, hidden_size).
        cusum_token_count (torch.Tensor): optional. Shape is (num_expert + 1).
        start_expert_id (int): begin expert, used in expert parrallelism.
        expert_size (int): expert size, used in expert parrallelism.

    Type:
        input: bfloat16, float16, float32.
        reduce_weight: float32.
        gather_ids: int32.
        residual: bfloat16, float16, float32.
        bias: bfloat16, float16, float32.
        cusum_token_count: int32.
        start_expert_id: int32.
        expert_size: int32.

    Return:
        Return output.
    '''
    return torch.ops.torch_mlu_ops.moe_combine_result(input,
                                                      reduce_weight,
                                                      gather_ids,
                                                      residual,
                                                      cusum_token_count,
                                                      start_expert_id,
                                                      expert_size,
                                                      bias)

def moe_quantize(x: torch.Tensor,
                 smooth: torch.Tensor,
                 zero: Optional[torch.Tensor] = None,
                 token_count: Optional[torch.Tensor] = None,
                 gather_index: Optional[torch.Tensor] = None,
                 gather_index_start_position: Optional[torch.Tensor] = None,
                 output: Optional[torch.Tensor] = None,
                 output_scale: Optional[torch.Tensor] = None,
                 dynamic_quant: bool = True
                ) -> Tuple[torch.Tensor]:
    """
    Apply quantization to the input tensor in MoE network.

    Math:
        if gather_idx and gather_index_start_position:
            start_pos = gather_index_start_position[0]
            input = input[gather_index[start_pos : start_pos + token_count.sum()]]
        elif gather_idx:
            input = input[gather_index]

        if token_count:
            input_list = input.split(token_count)
            group = token_count.size()
            result = []
            for i in range(group):
                result.append(input_list[i] * smooth[i])
            smoothed = concat(result, dim=0)
        else:
            smoothed = input * smooth

        if dynamic_quant:
            max, _ = smoothed.abs().max(dim=-1, keepdim=True)
            output_scale = max.to(torch.float) / 127.0
            quanted = (smoothed / output_scale).round().clamp(-128, 127).to(torch.int8)
            output = (quanted, output_scale)
        else:
            quanted = smoothed.round().clamp(-128, 127).to(torch.int8)
            output = (quanted)

    Args:
        input (torch.Tensor): The tensor to be quantized. Shape is (..., C) or (token_count.sum(), C).
        smooth (torch.Tensor): The smooth scale multipled to the input tensor. Shape is (group, C).
        zero (torch.Tensor): Not supported, must pass None.
        token_count (torch.Tensor): The tensor to separate input. Shape is (group).
        gather_index (torch.Tensor): The indices tensor to expand input. Shape is (expanded_tokens_num).
        gather_index_start_position (torch.Tensor): The tensor indicating the start position of gather_index. Shape is (1).
        output: (torch.Tensor): The tensor to store the output.
        output_scale: (torch.Tensor): The tensor to store output_scale.
        dynamic_quant: whether do dynamic quant.

    DataType:
        x: float, half or bfloat16.
        smooth: float.
        token_count: int32.
        gather_index: int32.
        gather_index_start_position: int32.
        dynamic_quant: bool.
        output: int8.
        output_scale: float.

    Return:
        Returns (output, output_scale) if dynamic_quant is True, otherwise returns output only.
    """
    output_shape = list(x.size())
    output_scale_shape = list(x.size()[:-1])
    if gather_index is not None:
        output_tokens = gather_index.size(0)
        output_shape[0] = output_tokens
        output_scale_shape[0] = output_tokens
    output = torch.empty(output_shape, dtype=torch.int8, device=x.device) if output is None else output
    output_scale = torch.empty(output_scale_shape, dtype=smooth.dtype, device=smooth.device) if output_scale is None and dynamic_quant else output_scale
    torch.ops.torch_mlu_ops.smooth_quant(x, smooth,
                                         output,
                                         output_scale,
                                         zero,
                                         token_count, gather_index,
                                         gather_index_start_position,
                                         'per_token', dynamic_quant)
    return (output, output_scale) if dynamic_quant else (output,)

def moe_active(input: torch.Tensor,
               act_mode: str,
               is_gated: bool,
               output: Optional[torch.Tensor] = None,
               bias: Optional[torch.Tensor] = None,
               cusum_token_count: Optional[torch.Tensor] = None,
               start_expert_id: int = 0,
               expert_size: int = 0) -> torch.Tensor:
    """
    Apply activation to the input tensor.

    Math:
        C = input.shape[-1]
        if bias:
            start_idx = 0
            for i in range(start_expert_id : start_expert_id + expert_size):
                deal_token_num = cusum_token_count[i+1] - cusum_token_count[i]
                input[start_idx:start_idx+deal_token_num] += bias[i]
                start_idx += deal_token_num
        if is_gated:
            output = active(input[..., :C//2]) * input[..., C//2:]
        else:
            output = active(input)

    Args:
        input (torch.Tensor): Shape is (..., C).
        act_mode (str): The activation mode, must be 'silu' or 'gelu'.
        is_gated (bool): If use gated activation.
        bias(torch.Tensor): optional, Shape is (expert_num, C).
        cusum_token_count(torch.Tensor): optional, Shape is (expert_num + 1).
        start_expert_id(int): optional, begin expert, used in expert parallelism.
        expert_size(int): optional, expert_size, used in expert parallelism.

    Type:
        input: float, half, bfloat16.
        act_mode: str.
        is_gated: bool.
        bias: same as input.
        cusum_token_count: int32.
        start_expert_id: int32.
        expert_size: int32.
        output: same as input.

    Return:
        if is_gated, output shape is (input.size()[:-1], C // 2),
        else, the same as input shape.
    """
    if output is None:
        if is_gated:
            output = torch.empty(input.size()[:-1] + (input.size()[-1] // 2,), dtype=input.dtype, device=input.device)
        else:
            output = torch.empty_like(input)
    torch.ops.torch_mlu_ops.active(input,
                                   output,
                                   bias,
                                   cusum_token_count,
                                   act_mode,
                                   is_gated,
                                   start_expert_id,
                                   expert_size)
    return output

def fused_rope(qkv: torch.Tensor,
               k_cache_hp: torch.Tensor,
               v_cache_hp: torch.Tensor,
               sin_cache: torch.Tensor,
               cos_cache: torch.Tensor,
               position_id: torch.Tensor,
               gamma: torch.Tensor,
               beta: torch.Tensor,
               k_cache_lp: Optional[torch.Tensor] = None,
               v_cache_lp: Optional[torch.Tensor] = None,
               cache_bs_id_hp: Optional[torch.Tensor] = None,
               cache_seq_offsets_hp: Optional[torch.Tensor] = None,
               cache_bs_id_lp: Optional[torch.Tensor] = None,
               cache_seq_offsets_lp: Optional[torch.Tensor] = None,
               k_scale_hp: Optional[torch.Tensor] = None,
               v_scale_hp: Optional[torch.Tensor] = None,
               k_scale_lp: Optional[torch.Tensor] = None,
               v_scale_lp: Optional[torch.Tensor] = None,
               slot_mapping_hp: Optional[torch.Tensor] = None,
               slot_mapping_lp: Optional[torch.Tensor] = None,
               eps: float = 1e-5) :
    """
    Perform query and key fold rope + key layernorm + (key, value perchannel quantize) +
    reshape key value to kv cache.

    Math:
        q = qkv[:, :, 0:head_num_q]
        k = qkv[:, :, head_num_q:head_num_q + head_num_k].clone
        v = qkv[:, :, head_num_q + head_num_k:].clone

        q = apply_rotary(q, sin_table, cos_table, position_ids)
        k = apply_rotary(k, sin_table, cos_table, position_ids)
        k = layernorm(k, gamma, beta)

        if (key_scale is not None and value is not None):
            k = quantize(k, key_scale)
            v = quantize(v, value_scale)

        if (slot_mapping is not None):
            reshape_paged_cache(key_cache, value_cache, k, v, slot_mapping)
        else:
            reshape_linear_cache(key_cache, value_cache, k, v, cache_bs_id, cache_seq_offset)

    Args:
        qkv (torch.Tensor): The qkv tensor. Shape is (batch, 1, head_num_q + head_num_k * 2, head_size).
        key_cache_hp (torch.Tensor): The high precision key cache tensor. Shape is (max_bs, head_num_k, max_decode_len, head_size) or
            (num_blocks, head_num_k, block_size, head_size).
        value_cache_hp (torch.Tensor): The high precision value cache tensor. Shape is (max_bs, head_num_k, max_decode_len, head_size) or
            (num_blocks, head_num_k, block_size, head_size).
        sin_table (torch.Tensor): The rotary sin table tensor. Shape is (rotary_seq, head_size).
        cos_table (torch.Tensor): The rotary cos table tensor. Shape is (rotary_seq, head_size).
        position_ids (torch.Tensor): The rotary seq_len offset of each batch. Shape is (batch).
        gamma (torch.Tensor): The layernorm gamma tensor. Shape is (head_size).
        beta (torch.Tensor): The layenorm beta tensor. Shape is (head_size).
        key_cache_lp (torch.Tensor): The low precision key cache tensor. Shape is (max_bs, head_num_k, max_decode_len, head_size) or
            (num_blocks, head_num_k, block_size, head_size / 2).
        value_cache_lp (torch.Tensor): The low precision value cache tensor. Shape is (max_bs, head_num_k, max_decode_len / 2, head_size) or
            (num_blocks, head_num_k, block_size / 2, head_size).
        cache_bs_id_hp (Torch.Tensor): The high precision cache batch offset tensor of each batch. Shape is (batch). This tensor
            can be None, that means cache_bs_id is (0, 1, 2, 3, 4, ...).
        cache_seq_offset_hp: (Torch.Tensor): The high precision cache seq_len offset tensor of each batch. Shape is (batch).
        cache_bs_id_lp (Torch.Tensor): The low precision cache batch offset tensor of each batch. Shape is (batch). This tensor
            can be None, that means cache_bs_id is (0, 1, 2, 3, 4, ...).
        cache_seq_offset_lp: (Torch.Tensor): The low precision cache seq_len offset tensor of each batch. Shape is (batch).
        key_scale_hp (Torch.Tensor): The high precision key scale tensor that quantize key to int8. Shape is (head_num_k, head_size).
        value_scale_hp (Torch.Tensor): The high precision value scale tensor that quantize value to int8. Shape is (head_num_k, head_size).
        key_scale_lp (Torch.Tensor): The low precision key scale tensor that quantize key to int8. Shape is (max_bs, head_num_k, group_num).
        value_scale_lp (Torch.Tensor): The low precision value scale tensor that quantize value to int8. Shape is (max_bs, head_num_k, group_num).
        slot_mapping_hp (Torch.Tensor): The slot_mapping tensor of high precision cache. Shape is (batch).
        slot_mapping_lp (Torch.Tensor): The slot_mapping tensor of low precision cache. Shape is (batch).
        eps (float): The layernorm eps param.

    Type:
        qkv: half, bfloat16.
        key_cache_hp: same as qkv if not quant kv else int8.
        value_cache_hp: same as qkv if not quant kv else int8.
        sin_table: same as qkv.
        cos_table: same as qkv.
        position_ids: int32.
        gamma: same as qkv.
        beta: same as qkv.
        key_cache_lp: int8(int4x2).
        value_cache_lp: int8(int4x2).
        cache_bs_id_hp: int32.
        cache_seq_offsets_hp: int32.
        cache_bs_id_lp: int32.
        cache_seq_offsets_lp: int32.
        key_scale_hp: float.
        value_scale_hp: float.
        key_scale_lp: float.
        value_scale_lp: float.
        slot_mapping_hp: int32.
        slot_mapping_lp: int32.

    Return:
        If both k_cache_lp and k_scale_lp are None, return a tuple of (qkv, k_cache_hp, v_cache_hp).
        If k_cache_lp exists but k_scale_lp is None, return a tuple of (qkv, k_cache_hp, v_cache_hp, k_cache_lp, v_cache_lp).
        If both k_cache_lp and k_scale_lp exist, return a tuple of (qkv, k_cache_hp, v_cache_hp, k_cache_lp, v_cache_lp, k_scale_lp, v_scale_lp).
    """
    torch.ops.torch_mlu_ops.fused_rope(qkv, k_cache_hp, v_cache_hp, k_cache_lp, v_cache_lp,
            sin_cache, cos_cache, position_id, gamma, beta, k_scale_hp, v_scale_hp,
            k_scale_lp, v_scale_lp, cache_bs_id_hp, cache_seq_offsets_hp,
            cache_bs_id_lp, cache_seq_offsets_lp, slot_mapping_hp, slot_mapping_lp, eps)
    out = (qkv, k_cache_hp, v_cache_hp)
    if k_cache_lp is not None:
        out += (k_cache_lp, v_cache_lp)
        if k_scale_lp is not None:
            out += (k_scale_lp, v_scale_lp)
    return out

def update_out_and_lse(out: torch.Tensor,
                       lse: torch.Tensor,
                       block_out: torch.Tensor,
                       block_lse: torch.Tensor,
                       seq_offsets: Optional[torch.Tensor] = None,
                       cu_seqs: Optional[torch.Tensor] = None,
                       block_cu_seqs: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor]:
    """
    Update out and log-sum-exp(lse) according to block out and block lse.

    Math:
        out = out - F.sigmoid(block_lse - lse) * (out - block_out)
        lse = lse - F.logsigmoid(lse - block_lse)
    Args:
        out (torch.Tensor): The output tensor. In pad mode, shape is (batch, max_seq_len, head_num, head_size).
                            In pack mode, shape is (total_seq_len, head_num, head_size).
        lse (torch.Tensor): The lse tensor. Shape is (batch, head_num, max_seq_len)
        block_out (torch.Tensor): The output tensor. In pad mode, shape is (batch, block_seq_len, head_num, head_size).
                                  In pack mode, shape is (total_block_seq_len, head_num, head_size).
        block_lse (torch.Tensor): The lse tensor. Shape is (batch, head_num, block_seq_len)
        seq_offsets (torch.Tensor): The seq offset of origin out and origin lse. Shape is (batch).
        cu_seqs (torch.Tensor): The cumulative sum of out seq_lens. Shape is (batch + 1).
        block_cu_seqs (torch.Tensor): The cumulative sum of block out seq_lens. Shape is (batch + 1).
    Type:
        out: half, bfloat16, float.
        lse: float.
        block_out: half, bfloat16, float.
        block_lse: float.
        seq_offsets: int32_t.
        cu_seqs: int32_t.
        block_cu_seqs: int32_t.
    Return:
        Support inplace outputs.
        Directly return the given out and lse.
    """
    torch.ops.torch_mlu_ops.update_out_and_lse(out, lse, block_out, block_lse,
                                               seq_offsets, cu_seqs, block_cu_seqs)
    return (out, lse)

def dequant_from_linear_cache(key: torch.Tensor,
                              value: Optional[torch.Tensor],
                              key_cache: torch.Tensor,
                              value_cache: Optional[torch.Tensor],
                              key_cache_quant_scale: torch.Tensor,
                              value_cache_quant_scale: Optional[torch.Tensor],
                              context_lengths: torch.Tensor,
                              max_context_len: int,
                              context_seq_offset: Optional[torch.Tensor],
                              cache_bs_id: Optional[torch.Tensor],
                              cache_seq_offset: Optional[torch.Tensor],
                              quant_mode: int = 0,
                              quant_bit: int = 8) -> None:
    """
    De-quantizes the key and value tensors from the provided linear cache and scale.

    This function is used during inference to convert the quantized key and value tensors stored in the cache back to
    their original floating-point format. It supports both 8-bit and 4-bit quantization and different quantization modes
    (per-channel or per-head).


    Math:
        Using the key de-quantization processes as an example:
        len_i = context_lengths[i]
        start_pos_i = sum(context_lentgths[:i]) if context_seq_offset is None else context_seq_offset[i]
        end_pos_i = start_pos_i + len_i
        key_i = key[start_pos_i:end_pos_i,...].transpose(1, 0)
        cache_pos = cache_bs_id[i]
        key_cache_i = key_cache[cache_pos,cache_seq_offset[i]:cache_seq_offset[i]+len[i],...]
        key_scale_i = key_scale[i,...] if quant_per_batch else key_scale
        if quant_mode == 0:
            key_scale_i = key_scale_i.reshape(head_num, -1, head_size)
        else:
            key_scale_i = key_scale_i.reshape(head_num, seq, -1)
        if quant_bit == 8:
            key_i[:] = ((key_cache_i.to(torch.float32) * key_scale_i).to(key.dtype)
        else:  # quant_bit == 4
            key_unpack_i = key_i.flatten()
            key_cache_flat_i = key_cache_i.flatten()
            key_unpack_i[0::2] = key_cache >> 4
            key_unpack_i[1::2] = key_cache << 4 >> 4
            key_i[:] = (key_unupack_i.reshape(k_i).to(torch.float32) * k_scale_quant_scale).to(key.dtype)

    Args:
        key (torch.Tensor): The key tensor with shape (total_seqlen, head_num, head_size).
        value (torch.Tensor, optional): The value tensor with shape (total_seqlen, head_num, head_size). Default is None.
        key_cache (torch.Tensor): The key cache tensor, shape depends on the quantization bit width,
                                  where max_batch should be greater than or equal to the actual batch size.
        - For 8-bit quantization: shape is (max_batch, head_num, cache_mem_len, head_size).
        - For 4-bit quantization: shape is (max_batch, head_num, cache_mem_len, head_size//2).
        value_cache (torch.Tensor, optional): The value cache tensor, shape depends on the quantization bit width. Default is None.
        - For 8-bit quantization: shape is (max_batch, head_num, cache_mem_len, head_size).
        - For 4-bit quantization: shape is (max_batch, head_num, cache_mem_len//2, head_size).
        key_cache_quant_scale (torch.Tensor): The quantization scale tensor for the key cache. Shape depends on the quantization mode.
        - For per-channel quantization: shape is (head_num, head_size).
        - For per-token quantization: shape is (max_batch, head_num, cache_mem_len).
        value_cache_quant_scale (torch.Tensor, optional): The quantization scale tensor for the value cache, same shape as key_cache_quant_scale. Default is None.
        context_lengths (torch.Tensor): The actual lengths of the input sequences within the current context, shape (batch).
        max_context_len (int32): The maximum length of all sequences in the current context, used to determine cache size.
        context_seq_offset (torch.Tensor, optional): The starting position offset of each input sequence in the context, where cache_mem_len should be greater than or equal to max_context_len.
        If not provided, it is calculated as the cumulative sum of context_lengths. Shape (batch). Default is None.
        cache_bs_id (torch.Tensor, optional): The batch index in the cache where the key and value tensors will be placed. Shape (batch). Default is None.
        cache_seq_offset (torch.Tensor, optional): The starting position offset of each input sequence in the cache. Shape (batch). Default is None.
        quant_mode (int, optional): Quantization mode: 0 for per-channel quantization (each channel uses a different quantization scale),
                                    1 for per-token quantization (each token uses a different quantization scale). Default is 0.
        quant_bit (int, optional): Quantization bit width: 8 for 8-bit quantization, 4 for 4-bit quantization. Default is 8.

    Type:
        key: half, bfloat16
        value: half, bfloat16
        key_cache: int8
        value_cache: int8
        key_cache_quant_scale: float32
        value_cache_quant_scale: float32
        context_lengths: int32
        max_context_len: int32
        context_seq_offset: int32
        cache_bs_id: int32
        cache_seq_offset: int32
        quant_mode: int32
        quant_bit: int32

    Return:
        None.
    """

    torch.ops.torch_mlu_ops.dequant_from_linear_cache(
        key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale,
        context_lengths, max_context_len, context_seq_offset, cache_bs_id, cache_seq_offset,
        quant_mode, quant_bit
    )

def dequant_from_paged_cache(key: torch.Tensor,
                             value: Optional[torch.Tensor],
                             key_cache: torch.Tensor,
                             value_cache: Optional[torch.Tensor],
                             key_cache_quant_scale: torch.Tensor,
                             value_cache_quant_scale: Optional[torch.Tensor],
                             context_lengths: torch.Tensor,
                             max_context_len: int,
                             context_seq_offset: Optional[torch.Tensor],
                             block_tables: torch.Tensor,
                             quant_mode: int = 0,
                             quant_bit: int = 8) -> None:
    """
    Dequantize `key_cache` and `value_cache` to `key` and `value` based on quantization parameters.

    Math:
        block_num = context_lengths[i] // block_size
        rem_token_num = context_lengths[i] % block_size
        # Index key_cache based on block_tables and context_lengths
        key_cache_i = torch.concat(
            [key_cache[block_tables[i, j], ...] for j in range(block_num)] +
            ([key_cache[block_tables[i, block_num], :, :rem_token_num, :]] if rem_token_num > 0 else [])
        )
        # Get context begin and end indices
        context_begin_i = torch.cumsum(context_lengths, dims=-1)[i - 1] if context_seq_offset is None else context_seq_offset[i]
        context_end_i = context_begin_i + context_lengths[i]

        # Slice key based on context_lengths and context_seq_offset
        key_i = key[context_begin_i:context_end_i, ...]

        # Determine the scale based on quant_mode
        if quant_mode == 0:
            key_scale_i = key_cache_quant_scale[:, None, :]
        elif quant_mode == 1:
            key_scale_i = torch.concat(
                [key_cache_quant_scale[block_tables[i, j], ..., None] for j in range(block_num)] +
                ([key_cache_quant_scale[block_tables[i, block_num], :, :rem_token_num, None]] if rem_token_num > 0 else [])
            )

        # Dequantize key_cache to key
        key_i[:] = (key_cache_i * key_scale_i).transpose(1, 0)

    Args:
        key (torch.Tensor): Output tensor for dequantized key. Shape is (total_seqlens, head_num, head_size).
        value (torch.Tensor): Output tensor for dequantized value. Shape is (total_seqlens, head_num, head_size).
        key_cache (torch.Tensor): Quantized key cache tensor. Shape is (total_blocks, head_num, block_size, head_size).
        value_cache (torch.Tensor): Quantized value cache tensor. Shape is (total_blocks, head_num, block_size, head_size).
        key_cache_quant_scale (torch.Tensor): Quantization scale for key cache. Shape is (head_num, head_size) for per-channel or (max_batch, head_num, block_size) for per-token.
        value_cache_quant_scale (torch.Tensor): Quantization scale for value cache. Shape is (head_num, head_size) for per-channel or (max_batch, head_num, block_size) for per-token.
        context_lengths (torch.Tensor): Lengths of contexts for each batch. Shape is (batch).
        max_context_len (int): Maximum context length.
        context_seq_offset (torch.Tensor): Optional sequence offset for context lengths. Shape is (batch).
        block_tables (torch.Tensor): Block tables for indexing. Shape is (batch, max_block_num).
        quant_mode (int): Quantization mode. 0 for per-channel, 1 for per-token. Default is 0.
        quant_bit (int): Quantization bit. Default is 8.

    Type:
        key: half, bfloat16
        value: half, bfloat16
        key_cache: int8
        value_cache: int8
        key_cache_quant_scale: float32
        value_cache_quant_scale: float32
        context_lengths: int32
        max_context_len: int
        context_seq_offset: int32
        block_tables: int32
        quant_mode: int
        quant_bit: int

    Return:
        None
    """

    torch.ops.torch_mlu_ops.dequant_from_paged_cache(
        key, value, key_cache, value_cache, key_cache_quant_scale, value_cache_quant_scale,
        context_lengths, max_context_len, context_seq_offset, block_tables, quant_mode, quant_bit
    )

# add dump_gese as decorator automatically for every function mentioned in __all__
add_gen_case_decorator(inspect.currentframe())
