.. _torch_mlu_ops:

torch_mlu_ops算子
==========================================

torch_mlu_ops算子为Torch-MLU-Ops设计和封装的PyTorch风格的算子API，能够以灵活的方式集成到社区或者客户的LLM推理引擎框架中，加速客户的LLM业务部署。

.. list-table:: torch_mlu_ops算子
    :widths: 2 4
    :header-rows: 1
    :class: longtable

    * - 序号
      - 算子名称

    * - 1
      - :ref:`torch_mlu_ops.active`
    * - 2
      - :ref:`torch_mlu_ops.apply_rotary`
    * - 3
      - :ref:`torch_mlu_ops.attention_project`
    * - 4
      - :ref:`torch_mlu_ops.batch_matmul`
    * - 5
      - :ref:`torch_mlu_ops.copy_blocks`
    * - 6
      - :ref:`torch_mlu_ops.copy_blocks_out_of_place`
    * - 7
      - :ref:`torch_mlu_ops.dequant_from_linear_cache`
    * - 8
      - :ref:`torch_mlu_ops.dequant_from_paged_cache`
    * - 9
      - :ref:`torch_mlu_ops.ffn`
    * - 10
      - :ref:`torch_mlu_ops.flash_attention`
    * - 11
      - :ref:`torch_mlu_ops.flash_attn_sq_mm_allreduce`
    * - 12
      - :ref:`torch_mlu_ops.fused_layer_norm`
    * - 13
      - :ref:`torch_mlu_ops.fused_moe`
    * - 14
      - :ref:`torch_mlu_ops.fused_norm_attention_project`
    * - 15
      - :ref:`torch_mlu_ops.fused_norm_residual_ffn`
    * - 16
      - :ref:`torch_mlu_ops.fused_rms_norm`
    * - 17
      - :ref:`torch_mlu_ops.fused_rope`
    * - 18
      - :ref:`torch_mlu_ops.group_gemm`
    * - 19
      - :ref:`torch_mlu_ops.matmul`
    * - 20
      - :ref:`torch_mlu_ops.matmul_allreduce`
    * - 21
      - :ref:`torch_mlu_ops.moe_active`
    * - 22
      - :ref:`torch_mlu_ops.moe_cast_gating`
    * - 23
      - :ref:`torch_mlu_ops.moe_combine_result`
    * - 24
      - :ref:`torch_mlu_ops.moe_expand_input`
    * - 25
      - :ref:`torch_mlu_ops.moe_gen_idx`
    * - 26
      - :ref:`torch_mlu_ops.moe_softmax_topk`
    * - 27
      - :ref:`torch_mlu_ops.offline_quant_to_linear_cache`
    * - 28
      - :ref:`torch_mlu_ops.offline_quant_to_paged_cache`
    * - 29
      - :ref:`torch_mlu_ops.preload`
    * - 30
      - :ref:`torch_mlu_ops.quant_to_linear_cache`
    * - 31
      - :ref:`torch_mlu_ops.quant_to_paged_cache`
    * - 32
      - :ref:`torch_mlu_ops.quantize`
    * - 33
      - :ref:`torch_mlu_ops.reshape_linear_cache`
    * - 34
      - :ref:`torch_mlu_ops.reshape_paged_cache`
    * - 35
      - :ref:`torch_mlu_ops.single_query_cached_kv_attn`
    * - 36
      - :ref:`torch_mlu_ops.single_query_mixed_cached_kv_attn`
    * - 37
      - :ref:`torch_mlu_ops.smooth_quant_matmul`
    * - 38
      - :ref:`torch_mlu_ops.smooth_quant_matmul_allreduce`
    * - 39
      - :ref:`torch_mlu_ops.swap_blocks`
    * - 40
      - :ref:`torch_mlu_ops.update_out_and_lse`
    * - 41
      - :ref:`torch_mlu_ops.weight_only_quant_matmul`


详细算子接口说明请参考以下章节。

.. _torch_mlu_ops.active:

torch_mlu_ops.active
------------------------------------------

.. code::

  torch_mlu_ops.active(input, act_mode, is_gated, active_coef)-> torch.Tensor

**功能描述**

激活算子。

**公式说明**

* 公式如下：

  .. math::

      \begin{array}{lcl}
      C = input.shape[-1]\\
      if is\_gated = True:\\
      output = active(input[..., :C//2]) * input[..., C//2:]\\
      else: \\
      output = active(input)
      \end{array}

* 公式中各参数描述如下：

  - ``input`` ：active算子的输入。
  - ``is_gated``：是否是gated结构。

* 公式中各函数描述如下：

  - ``active`` ：激活算子。

**参数说明**

- ``input``：Tensor类型，必选。输入维度 ``[..., C]``。 其中 ``C`` 为通道的个数，在最低维。
- ``act_mode``：string类型。可以是 ``silu`` ， ``gelu`` ， ``quick_gelu`` 或者 ``swish``。
- ``is_gated``：Bool类型。当 ``is_gated=true`` 时，为gated激活。
- ``active_coef``：FLOAT类型。默认值为1.0，只有在act_mode为swish的情况下才需要额外设置。

**返回值说明**

-  ``output``：激活后返回的tensor。

**规格限制**

- 数据类型支持：

  - input：FP32、FP16、BF16
  - output：FP32、FP16、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_add_bias_activation.py`` 。


.. _torch_mlu_ops.apply_rotary:

torch_mlu_ops.apply_rotary
------------------------------------------

.. code::

  torch_mlu_ops.apply_rotary(input,sin_cache,cos_cache,position_ids,cu_seqlens,interleaved,discrete,dynamic_ntk,max_seqlen)->Tensor

**功能描述**

计算Attention模块中将rotary embedding加在Q、K上的算子。

**公式说明**

请参考相关的GPT positon embedding文献：https://kexue.fm/archives/8265。该算子实现的主要功能为查表和加操作。


**参数说明**

- ``input``：Tensor类型，必选。Q或者K，pack模式下，维度 ``[total_seqlen, head_num, head_size]``，非pack模式下，维度 ``[batch_size, total_seqlen, head_num, head_size]``，其中 ``head_num`` 表示transformer模块中的头数，``head_size`` 表示每个头的维度。
- ``sin_cache``：Tensor类型，必选。输入计算rotary embedding的sin表，可能不是连续。当 ``dynamic_ntk=true`` 时，维度 ``[batch_size, rotary_seq_len, rotary_pos_emb_dim]``；当 ``dynamic_ntk=false`` 时，维度 ``[rotary_seq_len, rotary_pos_emb_dim]``。其中 ``rotary_seq_len`` 一般设置为大于等于 ``max_seqlen`` ，``rotary_pos_emb_dim`` 在full rotary情况下，一般等于 ``hidden_size``，在partial rotary情况下，需要根据网络进行配置。
- ``cos_cache``：Tensor类型，必选。输入计算rotary embedding的cos表，可能不是连续。当 ``dynamic_ntk=true`` 时，维度 ``[batch_size, rotary_seq_len, rotary_pos_emb_dim]``，当 ``dynamic_ntk=false`` 时，维度 ``[rotary_seq_len, rotary_pos_emb_dim]``。
- ``position_ids``：Tensor类型，可选。当 ``discrete=true``，维度 ``[total_seqlen]`` 当 ``discrete=false``，维度 ``[batch_size]`` 。
- ``cu_seqlens``：Tensor类型，可选。在pack模式下，记录了input中每个batch的序列长度，维度 ``[batch + 1]``。表示当前batch的sequence长度与下一个batch的sequence长度的差值。如果不设置，则input中每个batch的序列长度为 ``max_seqlens``。
- ``interleaved``：Bool类型，必选。表示rotary embedding的方式。当 ``interleaved=true`` 时，rotary embedding的方式为 ``cross`` ，当 ``interleaved=false`` 时，rotary embedding的方式为 ``fold`` 。
- ``discrete``：Bool类型，必选。discrete表示position_ids是否是离散的。如果是离散的，需要给出每个token的完整的 ``position_ids``，否则只需给出起始token的 ``position_ids`` ，算子内部根据起始 ``position_ids`` 和 ``seqlen`` 确定每个token的 ``position_ids``。默认情况下设置 ``discrete=False``。
- ``dynamic_ntk``：Bool类型，必选。``dynamic_ntk=true`` 时，表示所有的 ``batch`` 有不同的sin和cos值。当 ``dynamic_ntk=False`` 时，表示所有的 ``batch`` 有相同的sin和cos值。默认情况下设置 ``dynamic_ntk=False``。
- ``max_seqlen``：Int类型，必选。context阶段，在pad模式下，等于 ``pad_seqlen`` ，pack模式下等于 ``cu_seqlens[-1]``；generate阶段，等于 ``valid_token.max()`` + 解码帧数。

**返回值说明**

- ``output``：Tensor类型，必选。rotary embedding加到Q或者K上的输出，维度 ``[total_seqlen, head_num, head_size]``。其中 ``total_seqlen`` pad模式下等于 ``batch*pad_seqlen``， pack模式下等于 ``valid_seqlen`` 之和，即 ``cu_seqlens[-1]``。


**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - sin_cache：FP16、BF16
  - cos_cache：FP16、BF16
  - position_ids：Int32
  - cu_seqlens：Int32
  - output：FP16、BF16

**备注**

- 该算子中sin_cache/cos_cache的维度及生成方式和vLLM存在差异，具体表现在：

  - rope维度：vLLM中sin_cache/cos_cache的rope的维度是 ``rotary_pos_emb_dim / 2``, torch_mlu_ops中sin_cache/cos_cache的rope的维度是完整的rotary_pos_emb_dim；
  - 生成方式：

    * vLLM中生成fold/cross旋转角度公式 : ``1.0 / (base ** (torch.arange(0, rotary_pos_emb_dim, 2) / rotary_pos_emb_dim))``；
    * torch_mlu_ops中生成fold模式旋转角度公式  : ``1.0 / (base ** (torch.arange(0, rotary_pos_emb_dim, 1) % half_rotary_pos_emb_dim / rotary_pos_emb_dim))``, ``half_rotary_pos_emb_dim = rotary_pos_emb_dim // 2``；
    * torch_mlu_ops中生成cross模式旋转角度公式 : ``1.0 / (base ** (torch.arange(0, rotary_pos_emb_dim, 1) // 2 * 2 / rotary_pos_emb_dim))``。

- 差异性带来的影响：

  - 需要适配旋转角度的生成公式。
  - 无论sin_cache/cos_cache是何种维度形式, 都是针对特定硬件平台的设计的。


**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_apply_rotary.py`` 。


.. _torch_mlu_ops.attention_project:

torch_mlu_ops.attention_project
------------------------------------------

.. code::

  torch_mlu_ops.attention_project(input,weight,bias,residual,alpha,beta)->Tensor


**功能描述**

计算Attention模块中的Q、K、V算子。适用于GPT或者 ``hidden_size`` 大于2048的transformer网络。

**公式说明**

* 公式如下：

  .. math::

      \begin{array}{lcl}
      Q = fc(input, weight\_q) + bias\_q\\
      K = fc(input, weight\_k) + bias\_k\\
      V = fc(input, weight\_v) + bias\_v\\
      \end{array}

  当 ``residual`` 存在时：

    .. math::

      \begin{array}{lcl}
      Q_out = alpha*Q + beta*residual \\
      K_out = alpha*K + beta*residual\\
      V_out = alpha*V + beta*residual\\
      \end{array}

* 公式中各参数描述如下：

  - ``input`` ：attention_project算子的输入。
  - ``weight_q``、 ``weight_k`` 、 ``weight_v`` ：Query、Key、Value矩阵的权重。
  - ``bias_q``、 ``bias_k`` 、 ``bias_v`` ：Query、Key、Value矩阵的偏置。

* 公式中各函数描述如下：

  - ``fc`` ：全连接层（fully-connected）的缩写。

**参数说明**

- ``input``：Tensor类型，必选。``attention_project`` 的输入。维度 ``[batch_size, seq_lens, input_size]``。其中 ``batch_size`` 表示batch的大小，``seq_lens`` 表示sequence的长度，``input_size`` 表示输入的大小。
- ``weight``：Tensor类型，必选。``attention_project`` 的矩阵权重参数。Q、K、V权重参数被pack在一起。在MHA场景下，维度 ``[3*hidden_size, input_size]``。在MQA场景下，维度 ``[(head_num+2*head_num_kv)*head_size, input_size]``。其中 ``hidden_size`` 表示隐层的大小。
- ``bias``：Tensor类型，可选。``attention_project`` 的矩阵偏置参数。Q、K、V偏置参数被pack在一起。在MHA场景下，维度 ``[3*hidden_size]``。在MQA场景下，维度 ``[(head_num+2*head_num_kv)*head_size]``。其中 ``head_num_kv`` 表示MQA场景下kv的数量。
- ``residual``: Tensor类型，可选。``attention_project`` 的参差参数。维度与 ``input`` 保持一致。
- ``alpha`` ： float类型。默认值为1.0。当 ``residual`` 存在时，作为Q、K、V的系数。
- ``beta`` ： float类型。默认值为1.0。当 ``residual`` 存在时，作为 ``residual`` 的系数。


**返回值说明**

- output: Tensor类型，必选。attention_project的输出。Q、K、V被pack在一起。在MHA场景下，维度[batch_size, seq_lens, 3*hidden_size]。在MQA场景下，维度[batch_size, seq_lens, (head_num+2*head_num_kv)*head_size]。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - weight：FP16、BF16
  - bias：FP16、BF16
  - output：FP16、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_attn_proj.py`` 。


.. _torch_mlu_ops.batch_matmul:

torch_mlu_ops.batch_matmul
------------------------------------------

  torch_mlu_ops.batch_matmul(input, filter, residual, alpha, beta, a_scale, b_scale, trans_a, trans_b)-> torch.Tensor

**功能描述**

通用的批量矩阵乘法算子。

**公式说明**

* 批量matmul标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     mul_out = batch_matmul(input, filter)\\
     residual = alpha * mul_out + beta * residual\\
     output = residual
     \end{array}

* 公式中各参数描述如下：

  - ``input`` ：matmul的左矩阵。
  - ``filter``：matmul的右矩阵。
  - ``residual``：参差。
  - ``output`` ：计算最终的输出。

* 公式中各函数描述如下：

  - ``batch_matmul`` ：批量矩阵乘。

**参数说明**

- ``input``：Tensor类型，必选。矩阵乘的左矩阵。维度 ``[batch, mat_m, mat_k]``。
- ``filter``：Tensor类型，必选。矩阵乘的右矩阵。维度 ``[batch, mat_n, mat_k]``。
- ``residual``：Tensor类型，可选。参差。维度 ``[batch, mat_m, mat_n]``，如果有，计算结束后residual会被改变。
- ``alpha``：Float32类型，必选。矩阵乘法系数。
- ``beta``：Float32类型，必选。参差的系数。
- ``a_scale``：Float32类型，默认值为1.0。当input为INT8类型时对应的量化系数。
- ``b_scale``：Float32类型，默认值为1.0。当weight为INT8类型时对应的量化系数。
- ``trans_a``：Bool类型，默认值为False，标记input是否需要转置。
- ``trans_b``：Bool类型，默认值为True，标记weight是否需要转置。

**返回值说明**

- ``out``：Tensor类型，算子的输出，维度 ``[batch, mat_m, mat_n]``。

**规格限制**

- 数据类型支持：
  - input：FP16、BF16、FP32、INT8
  - filter：FP16、BF16、FP32、INT8
  - residual：FP16、BF16、FP32
  - out：FP16、BF16、FP32
  - alpha：FP32
  - beta：FP32
  - a_scale：FP32
  - b_scale：FP32
  - trans_a：BOOL
  - trans_b：BOOL

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_batch_matmul.py`` 。


.. _torch_mlu_ops.copy_blocks:

torch_mlu_ops.copy_blocks
----------------------------------------------------------------

.. code::

  torch_mlu_ops.copy_blocks(k_caches, v_caches, block_mapping)

**功能描述**

将 ``k_caches``、 ``v_caches`` 中源位置的BLOCK复制到指定位置。

**公式说明**

无。

**参数说明**

- ``k_caches``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, block_size, head_size>]``。其中 ``num_layers`` 表示layer数量， ``[num_heads, block_size, head_size]`` 整体大小作为每次需要copy的BLOCK大小。 ``num_blocks`` 表示BLOCK数量； ``block_size`` 表示输入序列中的特定块或子序列大小； ``num_heads`` 表示transformer模块中的头数； ``head_size`` 表示每个头的维度。
- ``v_caches``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, head_size, block_size>]``。其中每个参数含义与 ``k_caches`` 一致。
- ``block_mapping``：Map<int64, vector<int64>>类型，必选。第一个int64表示需要copy的源BLOCK位置（小于num_blocks），第二个int64表示将源BLOCK copy到的目标位置（小于num_blocks）， ``vector`` 表示一个源BLOCK可以copy到多个目标位置。

**返回值说明**

- 无，原位操作。

**规格限制**

- 数据类型支持：

  - k_caches：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - v_caches：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - block_mapping：INT64

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_copy_blocks.py`` 。


.. _torch_mlu_ops.copy_blocks_out_of_place:

torch_mlu_ops.copy_blocks_out_of_place
------------------------------------------------------------------------------------------------------------

.. code::

  torch_mlu_ops.copy_blocks_out_of_place(k_caches, v_caches, block_mapping)->Tuple(k_caches, v_caches)

**功能描述**

将 ``k_caches``、 ``v_caches`` 中源位置的BLOCK复制到指定位置。

**公式说明**

无。

**参数说明**

- ``k_caches``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, block_size, head_size>]``。其中 ``num_layers`` 表示layer数量， ``[num_heads, block_size, head_size]`` 整体大小作为每次需要copy的BLOCK大小。 ``num_blocks`` 表示BLOCK数量； ``block_size`` 表示输入序列中的特定块或子序列大小； ``num_heads`` 表示transformer模块中的头数； ``head_size`` 表示每个头的维度。
- ``v_caches``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, head_size, block_size>]``。其中每个参数含义与 ``k_caches`` 一致。
- ``block_mapping``：Map<int64, vector<int64>>类型，必选。第一个int64表示需要copy的源BLOCK位置（小于num_blocks），第二个int64表示将源BLOCK copy到的目标位置（小于num_blocks）， ``vector`` 表示一个源BLOCK可以copy到多个目标位置。

**返回值说明**

- k_caches, v_caches。

**规格限制**

- 数据类型支持：

  - k_caches：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - v_caches：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - block_mapping：INT64

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_copy_blocks.py`` 。


.. _torch_mlu_ops.dequant_from_linear_cache:

torch_mlu_ops.dequant_from_linear_cache
------------------------------------------

.. code::

  torch_mlu_ops.dequant_from_linear_cache(key,value,key_cache,value_cache,key_cache_scale,value_cache_scale,context_lengths,max_context_len,context_seq_offset,cache_bs_id,cache_seqlen_offset,quant_mode,quant_bit)

**功能描述**

将计算得到的K-V cache拼接到K-V上的反量化算子。


**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[total_seq_len, head_num, head_size]``。
- ``value``：Tensor类型，可选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[total_seq_len, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。当 ``quant_bit`` 为8时，维度为 ``[max_batch, head_num, cache_mem_len, head_size]``；当 ``quant_bit`` 为4时，维度为 ``[max_batch, head_num, cache_mem_len, head_size//2]``。其中 ``max_batch`` 要求大于等于 ``batch_size``，``cache_mem_len`` 为网络支持的最大sequences长度。
- ``value_cache``：Tensor类型，可选。Transformer模块中保存 ``Value cache`` 的Tensor。当 ``quant_bit`` 为8时，维度为 ``[max_batch, head_num, cache_mem_len, head_size]``；当 ``quant_bit`` 为4时，维度为 ``[max_batch, head_num, cache_mem_len//2, head_size]``。
- ``key_cache_scale``：Tensor类型，必选。Transformer模块中保存的 ``Key`` 的 ``scale`` 值。在context阶段，维度为 ``per-channel`` 量化 ``[head_num, head_size]`` 或者 ``per-head`` 量化 ``[max_batch, head_num, cache_mem_len]``。
- ``value_cache_scale``：Tensor类型，可选。Transformer模块中保存的 ``Value`` 的 ``scale`` 值。在context阶段，维度为 ``per-channel`` 量化 ``[head_num, head_size]`` 或者 ``per-head`` 量化 ``[max_batch, head_num, cache_mem_len]``。
- ``context_lengths``：Tensor类型，必选。为每个batch的sequences length长度。
- ``max_context_len``：Int类型，必选。在context阶段，表示最大的sequence length。
- ``context_seq_offset``：Tensor类型，可选。表示context的seq的累积偏移，需要保证每个batch的sequences不重合，维度 ``[batch_size]``。如果为空，则默认为context_lens的累加和。
- ``cache_bs_id``：Tensor类型，可选。表示cache的batch的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 1, 2 ... batch - 1}。
- ``cache_seq_offset``：Tensor类型，可选。表示cache的seq_lens的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 0, 0 ... 0}。
- ``quant_mode``: Int类型，可选。表示量化的模式，可选0（表示per-channel）或者1（表示per-head），默认值为0。
- ``quant_bit``：Int类型，可选。表示量化结果的位宽，可选8（表示int8）或4（表示int4），默认值为8。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16
  - value：FP16、BF16
  - key_cache：INT8
  - value_cache：INT8
  - key_cache_scale： FP32
  - value_cache_scale： FP32
  - context_lengths：INT32
  - max_context_len：INT32
  - context_seq_offset：INT32
  - cache_bs_id： INT32
  - cache_seqlen_offset： INT32
  - quant_mode: INT32
  - quant_bit: INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_dequant_from_linear_cache.py`` 。


.. _torch_mlu_ops.dequant_from_paged_cache:

torch_mlu_ops.dequant_from_paged_cache
------------------------------------------

.. code::

  torch_mlu_ops.dequant_from_paged_cache(key,value,key_cache,value_cache,key_cache_scale,value_cache_scale,context_lengths,max_context_len,context_seq_offset,block_tables,quant_mode,quant_bit)

**功能描述**

将计算得到的K-V cache（Paged模式）拼接到K-V上的反量化算子。


**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[total_seq_len, head_num, head_size]``。
- ``value``：Tensor类型，可选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[total_seq_len, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。当 ``quant_bit`` 为8时，维度为 ``[block_num, head_num, block_size, head_size]``。
- ``value_cache``：Tensor类型，可选。Transformer模块中保存 ``Value cache`` 的Tensor。当 ``quant_bit`` 为8时，维度为 ``[block_num, head_num, block_size, head_size]``。
- ``key_cache_scale``：Tensor类型，必选。Transformer模块中保存的 ``Key`` 的 ``scale`` 值。在context阶段，维度为 ``per-channel`` 量化 ``[head_num, head_size]`` 或者 ``per-head`` 量化 ``[block_num, head_num, block_size]``。
- ``value_cache_scale``：Tensor类型，可选。Transformer模块中保存的 ``Value`` 的 ``scale`` 值。在context阶段，维度为 ``per-channel`` 量化 ``[head_num, head_size]`` 或者 ``per-head`` 量化 ``[block_num, head_num, block_size]``。
- ``context_lengths``：Tensor类型，必选。为每个batch的sequences length长度。
- ``max_context_len``：Int类型，必选。在context阶段，表示最大的sequence length。
- ``context_seq_offset``：Tensor类型，可选。表示context的seq的累积偏移，需要保证每个batch的sequences不重合，维度 ``[batch_size]``。如果为空，则默认为context_lens的累加和。
- ``block_tables``：Tensor类型，可选。表示cache所在的block序号，维度 ``[batch_size, max_block_num]``。
- ``quant_mode``: Int类型，可选。表示量化的模式，可选0（表示per-channel）或者1（表示per-head），默认值为0。
- ``quant_bit``：Int类型，可选。表示量化结果的位宽，可选8（表示int8），默认值为8。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16
  - value：FP16、BF16
  - key_cache：INT8
  - value_cache：INT8
  - key_cache_scale： FP32
  - value_cache_scale： FP32
  - context_lengths：INT32
  - max_context_len：INT32
  - context_seq_offset：INT32
  - block_tables: INT32
  - quant_mode: INT32
  - quant_bit: INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_dequant_from_paged_cache.py`` 。


.. _torch_mlu_ops.ffn:

torch_mlu_ops.ffn
------------------------------------------

.. code::

  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,act_mode)->Tensor

**功能描述**

计算transformer中的FFN（FeedForward）前馈神经网络模块的算子。适用于GPT或者 ``hidden_size`` 大于2048的transformer网络。

**公式说明**

* FFN标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     fc\_output = fc(input, up\_fc\_weight) + up\_fc\_bias\\
     active\_output = activation(fc\_output)\\
     basic\_output = fc(active\_output, down\_proj\_weight) + down\_proj\_bias
     \end{array}

* Gated FFN标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     fc\_output = fc(input, up\_fc\_weight) + up\_fc\_bias\\
     active\_output = activation(fc\_output)\\
     gated_fc\_output = fc(input, gate\_up\_proj\_weight) + gate\_up\_proj\_bias\\
     elementwise\_product\_output = elementwise\_product(active\_output, gated\_fc\_output)\\
     gated\_output = fc(elementwise\_product\_output, down\_proj\_weight) + down\_proj\_bias\\
     \end{array}

* 公式中各参数描述如下：

  - ``input`` ：FFN模块的输入。
  - ``fc_output`` ：FFN模块经过升维后的输出。
  - ``active_output`` ：FFN模块经过激活后的输出。
  - ``basic_output`` ：标准FFN模块经过降维后的输出。
  - ``gated_fc_output`` ：gated FFN旁路经过升维后的输出。
  - ``elementwise_product_output`` ：gated FFN主路和旁路经过elementwise product后的输出。
  - ``gated_output`` ：gated FFN模块经过降维后的输出。


* 公式中各函数描述如下：

  - ``fc`` ：全连接层（fully-connected）的缩写。
  - ``activation`` ：用户可配的激活函数，可选的激活函数："relu"、"gelu"、"silu"。
  - ``elementwise_product`` ：点积。

**参数说明**

- ``input``：Tensor类型，必选。FFN算子的输入。维度[batch_size, seq_lens, hidden_size]。
- ``up_fc_weight``：Tensor类型，必选。FFN算子中升维矩阵权重参数，维度[ffn_inner_size, hidden_size]。
- ``up_fc_bias``：Tensor类型，可选。FFN算子中升维矩阵偏置参数，维度[ffn_inner_size]。其中ffn_inner_size表示FFN中间层的维度。
- ``down_proj_weight``：Tensor类型，必选。FFN算子中降维矩阵权重参数，维度[hidden_size, ffn_inner_size]。
- ``down_proj_bias``：Tensor类型，可选。FFN算子中降维矩阵偏置参数，维度[hidden_size]。
- ``gate_up_proj_weight``：Tensor类型，可选。Gated FFN算子中升维矩阵权重参数，维度[hidden_size, ffn_inner_size]。
- ``gate_up_proj_bias``：Tensor类型，可选。Gated FFN算子中升维矩阵偏置参数，维度[hidden_size]。
- ``act_mode``：string类型，必选。FFN的激化函数类型，可选的激活函数："relu"、"gelu"、"silu"、"none"。

**返回值说明**

- output: Tensor类型，必选。ffn的输出。维度[batch_size, seq_lens, hidden_size]。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - up_fc_weight：FP16、BF16
  - up_fc_bias：FP16、BF16
  - down_proj_weight：FP16、BF16
  - down_proj_bias: FP16、BF16
  - gate_up_proj_weight: FP16、BF16
  - gate_up_proj_bias: FP16、BF16
  - output: FP16、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_ffn.py`` 。


.. _torch_mlu_ops.flash_attention:

torch_mlu_ops.flash_attention
------------------------------------------

.. code::

  torch_mlu_ops.flash_attention(q,k,v,out,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,compute_dtype,return_lse)

**功能描述**

Context阶段的Flash Attention算子。

**公式说明**

* flash attention标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     scores = batch\_matmul(q, k, transpose\_b=True) * softmax\_scale\\
     scores += ALiBi + attn\_bias + causal\_mask + window\_mask\\
     output\_lse = logsumexp(scores)\\
     attention = softmax(scores)\\
     output = batch\_matmul(attention, v)\\
     \end{array}

* 公式中各参数描述如下：

  - ``q/k/v`` ：flash attention模块的输入。
  - ``attn_bias``：flash attention模块的偏置。
  - ``ALiBi``：flash attention模块在线计算的ALiBi位置编码矩阵。
  - ``causal_mask``：如果是因果attention的，则需要添加mask。
  - ``window_mask``：如果设置了window的尺寸，则需要添加mask。
  - ``scores`` ：flash attention模块第一个批量矩阵乘的输出。
  - ``output_lse``：flash_attention的lse输出。
  - ``attention`` ：flash attention模块经过softmax后的输出。
  - ``output`` ：flash attention模块第二个批量矩阵乘的输出，也是模块的最终输出。

* 公式中各函数描述如下：

  - ``batch_matmul`` ：批量矩阵乘。
  - ``softmax`` ：归一化指数函数。
  - ``logsumexp`` ：对数求和指数函数。

**参数说明**

- ``q``：Tensor类型，必选。flash attention算子的输入。维度 ``[total_q, q_head_num, head_size_qk]``。
- ``k``：Tensor类型，必选。flash attention算子的输入。当block_tables为空时，维度 ``[total_k, kv_head_num, head_size_qk]``， 否则维度 ``[num_blocks, kv_head_num, block_size, head_size_qk]``。
- ``k``：Tensor类型，必选。flash attention算子的输入。当block_tables为空时，维度 ``[total_k, kv_head_num, head_size_v]``， 否则维度 ``[num_blocks, kv_head_num, block_size, head_size_v]``。
- ``out``：Tensor类型，可选。flash attention算子的原位输入。维度同input。
- ``cu_seq_lens``：Tensor类型，可选。当输入q/k/v为pack模式时候记录每个batch的seq的位置，维度 ``[batch+1]``。
- ``alibi_slope``：Tensor类型，可选。是一种相对位置编码方法，维度 ``[batch, q_head_num]``或者 ``[q_head_num]``，类型必须是float32。
- ``attn_bias``：Tensor类型，可选。attention的偏置，维度 ``[batch, max_seq_len, max_seq_len]`` 或者 ``[batch, q_head_num, max_seq_len, max_seq_len]``。
- ``k_cache_quant_scale``：Tensor类型，预留参数，必须为空。
- ``v_cache_quant_scale``：Tensor类型，预留参数，必须为空。
- ``block_tables``：Tensor类型，可选。block的id与kv_cache之间的映射关系，为空时表示非paged模式。维度 ``[bs, max_num_blocks_per_seq]``，max_num_blocks_per_seq为1时block_size=max(seq_len_kv)。
- ``max_seq_len``：Int32类型，必选。算子输入的最大seq长度。
- ``softmax_scale``：Float32类型，必选。softmax之前的系数。
- ``is_causal``：Bool类型，标记是否是因果的，如果为因果的则计算过程中需要使用 ``causal_mask``。
- ``window_size_left``：Int32类型, 左侧window size大小，无限制则设置为-1。实际参与计算的key和value长度为 ``window_size_left + 1``。
- ``window_size_right``：Int32类型，右侧window size大小，无限制则设置为-1。实际参与计算的key和value长度为 ``window_size_right + 1``。
- ``return_lse``：Bool类型，标记是否返回lse tensor。


**返回值说明**

- ``output``：Tensor类型，可选。flash attention的输出，当输入的out不为空时无此输出。维度 ``[total_q, q_head_num, head_size_v]``。
- ``output_lse``：Tensor类型，可选。flash attention的输出，当return_lse为true时有此输出，维度 ``[batch, head, max_seq_len_q]``。


**规格限制**

- 数据类型支持：

  - q/k/v：FP16、BF16、FP32
  - cu_seq_lens：INT32
  - alibi_slope：FP32
  - attn_bias：FP16、BF16、FP32
  - block_tables：INT32
  - max_seq_len：INT32
  - softmax_scale：FP32
  - is_causal：BOOL
  - window_size_left：INT32
  - window_size_right：INT32
  - output：FP16、BF16、FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_flash_attention.py``。


.. _torch_mlu_ops.flash_attn_sq_mm_allreduce:

torch_mlu_ops.flash_attn_sq_mm_allreduce
------------------------------------------

.. code::

  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,
                                           compute_dtype, block_seq)

**功能描述**

Context阶段的Flash Attention、smooth_quant matmul与allreduce通信计算融合算子。

**公式说明**

无

**参数说明**

- ``cncl_comm`` ：Int类型，cnclComm，用于归约通信。
- ``q``：Tensor类型，必选。flash attention算子的输入。维度 ``[total_q, q_head_num, head_size]``。
- ``k/v``：Tensor类型，必选。flash attention算子的输入。维度 ``[total_k, kv_head_num, head_size]``。
- ``cu_seq_lens_q``：Tensor类型，可选。当输入q为pack模式时候记录每个batch的seq的位置，维度 ``[batch + 1]``。
- ``cu_seq_lens_kv``：Tensor类型，可选。当输入k/v为pack模式时候记录每个batch的seq的位置，维度 ``[batch + 1]``。
- ``alibi_slope``：Tensor类型，可选。是一种相对位置编码方法，维度 ``[batch, q_head_num]`` 或者 ``[q_head_num]``，类型必须是float32。
- ``attn_bias``：Tensor类型，可选。attention的偏置，维度 ``[batch, max_seq_len, max_seq_len]`` 或者 ``[batch, q_head_num, max_seq_len, max_seq_len]``。
- ``smooth``：Tensor类型，必选。平滑系数, 维度 ``[head_num_q * head_size]``。
- ``weight``：Tensor类型，必选。, smooth_quant matmul的权重矩阵，维度 ``[head_num_q * head_size, head_num_q * head_size]``。
- ``weight_scale``：Tensor类型，必选。, smooth_quant matmul的权重矩阵的量化参数，维度 ``[head_num_q * head_size]``。
- ``bias``：Tensor类型，可选。, smooth_quant matmul的偏置，维度 ``[head_num_q * head_size]``。
- ``max_seq_len_q``：Int32类型，必选。q的最大seq长度。
- ``max_seq_len_kv``：Int32类型，必选。k/v的最大seq长度。
- ``softmax_scale``：Float32类型，必选。softmax之前的系数。
- ``is_causal``：Bool类型，标记是否是因果的，如果为因果的则计算过程中需要使用 ``causal_mask``。
- ``window_size_left``：Int32类型, 左侧window size大小，无限制则设置为-1。
- ``window_size_right``：Int32类型, 右侧window size大小，无限制则设置为-1。
- ``compute_dtype``：torch类型，在计算softmax时使用。
- ``block_seq``: Int类型，保留字段，设置0即可。


**返回值说明**

- ``output``：Tensor类型，必选。输出，维度 ``[total_q, q_head_num * head_size]``。


**规格限制**

- 数据类型支持：

  - cncl_comm: INT64
  - q/k/v：FP16、BF16、FP32
  - cu_seq_lens_q：INT32
  - cu_seq_lens_kv：INT32
  - alibi_slope：FP32
  - attn_bias：FP16、BF16、FP32
  - smooth：FP32
  - weight：INT8
  - weight_scale：FP32
  - bias：FP16、BF16、FP32
  - max_seq_len_q：INT32
  - max_seq_len_kv：INT32
  - softmax_scale：FP32
  - is_causal：BOOL
  - window_size_left：INT32
  - window_size_right：INT32
  - block_seq：INT32
  - output：FP16、BF16、FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/allreduce_case/test_flash_attn_sq_mm_allreduce.py``。


.. _torch_mlu_ops.fused_layer_norm:

torch_mlu_ops.fused_layer_norm
------------------------------------------

.. code::

  torch_mlu_ops.fused_layer_norm(input,residual,gamma,beta,bias,eps,store_output_before_norm,quant_scale,out,dynamic_quant)


**功能描述**

layernorm归一化算子。

**公式说明**

* layernorm标准结构公式如下：

  .. math::

    \begin{aligned}
    output = weight \times \frac{input - Mean(input)}{\sqrt{Variance(input) + eps}} + bias\\
    output = Quantize(output, quant_scale)
    \end{aligned}

* 公式中各参数描述如下：

  - ``input`` layernorm的输入。
  - ``Mean()`` 表示对input取均值。
  - ``Variance()`` 表示对input取方值。
  - ``Quantize()`` 表示对输出做量化。
  - ``weight`` 表示layernorm的 ``weight`` 参数。
  - ``bias`` 表示layernorm的 ``bias`` 参数。
  - ``eps`` 表示layernorm eps的值，一个极小的正数，避免除零，一般默认值设定为0.00001。
  - ``output`` layernorm的输出。
  - ``quant_scale`` 量化参数。

**参数说明**

- ``input``：Tensor类型，必选。layernorm算子的输入。维度 ``[batch，seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``residual``：Tensor类型，可选。layernorm算子的残差输入。在pad模式下，维度 ``[batch，seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``gamma``：Tensor类型，可选。layernorm算子的 ``weight`` 。维度 ``[hidden_size]``。
- ``beta``：Tensor类型，可选。layernorm算子的 ``bias`` 。维度 ``[hidden_size]``。
- ``bias``：Tensor类型，可选。layernorm算子的残差的偏置输入。维度 ``[hidden_size]``。
- ``eps``：Float类型，必选。layernorm算子的 ``eps`` 的值，一般设定为0.00001。
- ``store_output_before_norm``：Bool类型，必选。layernorm算子是否输出 ``input+residual`` 的值。当为true时，第二个返回值输出 ``input+residual``，当为false时，无第二个返回值。
- ``quant_scale``：Float类型，可选。量化参数。维度 ``[hidden_size]``。当 ``dynamic_quant`` 为False时表示离线量化参数，为True时表示在线量化时的平滑因子。当该参数不为空时，输入输出必须连续。
- ``out``：Tensor类型，可选。layernorm算子的原位输出。维度同input，支持batch维度和seq_lens维度的stride。
- ``dynamic_quant``：Bool类型，为False且 ``quant_scale`` 不为空时表示离线per-channel量化输出，否则为在线动态per-token量化输出。

**返回值说明**

- ``output``：Tensor类型，可选。layernorm算子的输出，当输入参数out不为空时无此输出。维度 ``[batch_size，seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``residual_output``：Tensor类型，可选。layernorm算子的residual输出。当store_output_before_norm为true时，返回值输出 ``input+residual``，当store_output_before_norm为false时，无此返回值。在pad模式下，维度 ``[batch_size*seq_lens, hidden_size]``；在pack模式下，维度 ``[total_seqlen, hidden_size]``。
- ``smooth_quant_scale``：Tensor类型，可选。维度 ``[batch_size*seq_lens]`` 。当 ``dynamic_quant`` 为True时输出的量化scale。


**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - residual：FP16、BF16
  - beta：FP16、BF16
  - gamma：FP16、BF16
  - bias：FP16、BF16
  - eps：Float32
  - store_output_before_norm：Bool
  - quant_scale：Float
  - dynamic_quant：Bool
  - output：FP16、BF16、int8
  - residual_output： FP16、BF16
  - smooth_quant_scale： FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_fusenorm.py`` 。

.. _torch_mlu_ops.fused_moe:

torch_mlu_ops.fused_moe
------------------------------------------

.. code::

  torch_mlu_ops.fused_moe(input, router_logit, w1, w2, bias1, bias2,
                   residual, input_smooth, act_smooth, w1_scale,
                   w2_scale, topk, renormalize, gated, act_mode,
                   start_expert_id, block_n, cncl_comm, w1_quant_flag, w2_quant_flag)

**功能描述**

mixtral of experts, FFN（FeedForward）前馈神经网络模块的平替层，支持通信计算融合模式。

**公式说明**

无

**参数说明**

- ``input``：Tensor类型，必选。MOEFFN算子的输入。维度 ``[batch_size, seq_lens, hidden_size]`` 或 ``[total_seq, hidden_size]``。
- ``router_logit``：Tensor类型，必选。MOEFFN算子中expert的得分。维度 ``[batch_size, seq_lens, expert_num]`` 或 ``[batch_size, seq_lens, expert_num]``。
- ``w1``：Tensor类型，必选。MOEFFN算子中升维矩阵权重参数。维度 ``[expert_size, ffn_inner_size*(1+gated), hidden_size]`` 或 ``[expert_size, ffn_inner_size*(1+gated), hidden_size / 2]``。
- ``w2``：Tensor类型，必选。MOEFFN算子中降维矩阵权重参数。维度 ``[expert_size, hidden_size, ffn_inner_size]`` 或 ``[sub_expert_size, hidden_size, block, ffn_inner_size]`` 或 ``[expert_size, hidden_size, ffn_inner_size / 2]``。
- ``bias1``：Tensor类型，可选。MOEFFN算子中升维矩阵偏置参数。维度 ``[expert_size, ffn_inner_size*(1+gated)]``。
- ``bias2``：Tensor类型，可选。MOEFFN算子中降维矩阵偏置参数。维度 ``[expert_size, hidden_size]``。
- ``residual``：Tensor类型，可选。MOEFFN算子的残差。维度 ``[batch_size, seq_lens, hidden_size]`` 或 ``[total_seq, hidden_size]``。
- ``input_smooth``：Tensor类型，可选。MOEFFN算子smooth_quant模式时input的平滑系数。维度 ``[hidden_size]``。
- ``act_smooth``：Tensor类型，可选。MOEFFN算子smooth_quant模式时激活操作之后的平滑系数。维度 ``[ffn_inner_size]``。
- ``w1_scale``：Tensor类型，可选。MOEFFN算子量化模式时升维权重的量化参数。维度 ``[expert_size, ffn_inner_size*(1+gated)]`` 或 ``[w1_quant_group, expert_size, ffn_inner_size*(1+gated)]``。
- ``w2_scale``：Tensor类型，可选。MOEFFN算子量化模式时降维权重的量化参数。维度 ``[expert_size, hidden_size]`` 或 ``[w2_quant_group, expert_size, hidden_size]``。
- ``topk``：Int类型，必选。MOEFFN算子中，从每个token的expert_logits中选择topk个。
- ``renormalize``：Bool类型，必选。MOEFFN算子中softmax结果后是否需要归一化。
- ``gated``：Bool类型，必选。MOEFFN层中是否有gated结构。
- ``act_mode``：string类型，必选。MOEFFN的激化函数类型，可选的激活函数：``gelu``、``silu``。
- ``start_expert_id``：Int类型，可选。用于表明在专家并行模式下专家的起始位置。
- ``block_n``：Int类型，可选。用于表明在通算融合模式下group_gemm和allreduce的拆分次数。
- ``cncl_comm``：Int类型，可选。用于表明在通算融合模式下的通信句柄。
- ``w1_quant_flag``：List类型，可选。MOEFFN算子w4w8混合量化模式时标记quant_group的bit位，数值必须是4或8，其它值会引发错误。长度 ``[expert_size * w1_quant_group]``。
- ``w2_quant_flag``：List类型，可选。MOEFFN算子w4w8混合量化模式时标记quant_group的bit位，数值必须是4或8，其它值会引发错误。长度 ``[expert_size * w2_quant_group]``。


**返回值说明**

- ``output``: Tensor类型。 维度 ``[batch_size, seq_lens, hidden_size]`` 或 ``[total_seq, hidden_size]``。

**规格限制**

- 参数限制：

  - input_smooth/act_smooth/w1_scale/w2_scale: 只有smooth quant pertoken模式启用，input采用pertoken量化，权重采用per-channel量化。

- 数据类型支持：

  - input：HALF、FP32、BF16
  - router_logit：HALF、FP32、BF16
  - w1：HALF、FP32、BF16、INT8、INT8(INT4x2)
  - w2：HALF、FP32、BF16、INT8、INT8(INT4x2)
  - bias1：HALF、FP32、BF16
  - bias2：HALF、FP32、BF16
  - residual：HALF、FP32、BF16
  - input_smooth：FP32
  - act_smooth：FP32
  - w1_scale：FP32
  - w2_scale：FP32
  - start_expert_id: INT32
  - block_n: INT32
  - cncl_comm: INT64
  - output：HALF、FP32、BF16
  - w1_quant_flag: INT
  - w2_quant_flag: INT

**备注**

- w2维度是4D时，sub_expert_size * block 必须等于 expert_size, 并且要求ffn_inner_size是1024的整数倍。
- group量化时，w1_quant_group * w1_quant_group_size 必须等于 hidden_size，w2_quant_group * w2_quant_group_size 必须等于 ffn_inner_size。
- group量化时，w1_scale的形状必须是 (w1_quant_group, expert_size, ffn_inner_size*(1+gated))，w2_scale的形状必须是 (w2_quant_group, expert_size, hidden_size)。
- w1_quant_group_size与w2_quant_group_size表示group量化场景下一个group block的元素数量。必须是128的整数倍。
- w2维度是4D时不支持int4量化。
- 通算融合模式不支持int4量化。
- w1_quant_flag 与 w2_quant_flag存在时为w4w8混合量化模式，其元素只能是4或8。4表示该quant_group对应的权重是int4量化，8表示该quant_group对应的权重是int8量化。
- w4w8混合量化模式要求w1与w2必须是1D。

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe.py`` 和 ``torch_mlu_ops/tests/ops_pytest/allreduce_case/test_moe_all_reduce.py``。


.. _torch_mlu_ops.fused_norm_attention_project:

torch_mlu_ops.fused_norm_attention_project
------------------------------------------

.. code::

  torch_mlu_ops.fused_norm_attention_project(input,q_weight,q_bias,k_weight,k_bias,v_weight,v_bias,norm_weight,norm_bias,eps,out_layout,head_size,norm_out)->Tuple(Tensor)


**功能描述**

计算Attention模块中的Q、K、V算子。前融合 ``layernorm`` 的 ``attention_project``算子。 适用于GPT或者 ``hidden_size`` 小于2048的transformer网络。

**公式说明**

* 公式如下：

  .. math::

      \begin{array}{lcl}
      normed = layernorm(input, norm_weight, norm_bias, eps)
      Q = fc(normed, weight\_q) + bias\_q\\
      K = fc(normed, weight\_k) + bias\_k\\
      V = fc(normed, weight\_v) + bias\_v\\
      \end{array}


* 公式中各参数描述如下：

  - ``input`` ：fused_norm_attention_project算子的输入。
  - ``normed`` ： 经过 ``layernorm`` 计算后的输出。
  - ``weight_q``、 ``weight_k`` 、 ``weight_v`` 、 ``norm_weight`` ：Query、Key、Value矩阵的权重。
  - ``bias_q``、 ``bias_k`` 、 ``bias_v`` 、 ``nrom_bias`` ：Query、Key、Value矩阵的偏置。
  - ``eps`` ： layernorm中的 ``eps`` 参数。

* 公式中各函数描述如下：

  - ``fc`` ：全连接层（fully-connected）的缩写。
  - ``layernorm`` ： layernorm归一化。

**参数说明**

- ``input``：Tensor类型，必选。``fused_norm_attention_project`` 的输入。维度 ``[batch_size, seq_lens, input_size]``。其中 ``batch_size`` 表示batch的大小，``seq_lens`` 表示sequence的长度，``input_size`` 表示输入的大小。
- ``q_weight``：Tensor类型，必选。``fused_norm_attention_project`` 的Q矩阵权重参数。维度 ``[hidden_size, input_size]``。
- ``q_bias``：Tensor类型，可选。``fused_norm_attention_project` 的Q矩阵偏置参数。维度 ``[hidden_size]``。
- ``k_weight``：Tensor类型，必选。``fused_norm_attention_project`` 的K矩阵权重参数。维度 ``[hidden_size, input_size]``。
- ``k_bias``：Tensor类型，可选。``fused_norm_attention_project` 的K矩阵偏置参数。维度 ``[hidden_size]``。
- ``v_weight``：Tensor类型，必选。``fused_norm_attention_project`` 的V矩阵权重参数。维度 ``[hidden_size, input_size]``。
- ``v_bias``：Tensor类型，可选。``fused_norm_attention_project` 的V矩阵偏置参数。维度 ``[hidden_size]``。
- ``norm_weight``：Tensor类型，可选。``fused_norm_attention_project`` 的layernorm权重参数。维度 ``[input_size]``。
- ``norm_bias``：Tensor类型，可选。``fused_norm_attention_project` 的Vlayernorm偏置参数。维度 ``[input_size]``。
- ``eps`` ： float类型。layernorm的 ``eps`` 参数。
- ``out_layout`` ： String类型。允许的layout为 ``nthc`` 或者 ``nhtc``。
- ``head_size`` : String类型。 ``output`` 的最后一维，当 ``out_layout`` 为 ``nhtc`` 。
- ``norm_out`` : Bool类型。 当等于 ``True`` 时，返回 ``layernorm`` 的输出。


**返回值说明**

- output: Tensor类型，必选。fused_norm_attention_project的输出。根据入参的条件，返回不同的结果。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - q_weight：FP16、BF16
  - q_bias：FP16、BF16
  - k_weight：FP16、BF16
  - k_bias：FP16、BF16
  - v_weight：FP16、BF16
  - v_bias：FP16、BF16
  - norm_weight：FP16、BF16
  - norm_bias：FP16、BF16
  - output：FP16、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_fused_attn_proj.py`` 。


.. _torch_mlu_ops.fused_norm_residual_ffn:

torch_mlu_ops.fused_norm_residual_ffn
------------------------------------------

.. code::

  torch_mlu_ops.fused_norm_residual_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,eps,act_mode,residual_is,alhpa,beta)->Tensor

**功能描述**

计算transformer中的FFN（FeedForward）前馈神经网络模块的算子。前融合layernorm算子，适用于GPT或者 ``hidden_size`` 小于2048的transformer网络。

**公式说明**

* 前融合layernorm算子的FFN标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     normed = layernorm(input, norm_weight, norm_bias, eps)
     fc\_output = fc(normed, up\_fc\_weight) + up\_fc\_bias\\
     active\_output = activation(fc\_output)\\
     basic\_output = fc(active\_output, down\_proj\_weight) + down\_proj\_bias
     \end{array}

* 前融合layernorm算子的Gated FFN标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     normed = layernorm(input, norm_weight, norm_bias, eps)
     fc\_output = fc(normed, up\_fc\_weight) + up\_fc\_bias\\
     active\_output = activation(fc\_output)\\
     gated_fc\_output = fc(input, gate\_up\_proj\_weight) + gate\_up\_proj\_bias\\
     elementwise\_product\_output = elementwise\_product(active\_output, gated\_fc\_output)\\
     gated\_output = fc(elementwise\_product\_output, down\_proj\_weight) + down\_proj\_bias\\
     \end{array}

* 公式中各参数描述如下：

  - ``input`` ：FFN模块的输入。
  - ``normed`` ： 经过 ``layernorm`` 计算后的输出。
  - ``fc_output`` ：FFN模块经过升维后的输出。
  - ``active_output`` ：FFN模块经过激活后的输出。
  - ``basic_output`` ：标准FFN模块经过降维后的输出。
  - ``gated_fc_output`` ：gated FFN旁路经过升维后的输出。
  - ``elementwise_product_output`` ：gated FFN主路和旁路经过elementwise product后的输出。
  - ``gated_output`` ：gated FFN模块经过降维后的输出。


* 公式中各函数描述如下：

  - ``fc`` ：全连接层（fully-connected）的缩写。
  - ``activation`` ：用户可配的激活函数，可选的激活函数："relu"、"gelu"、"silu"。
  - ``elementwise_product`` ：点积。

**参数说明**

- ``input``：Tensor类型，必选。 ``fused_norm_residual_ffn`` 算子的输入。维度 ``[batch_size, seq_lens, hidden_size]`` 。
- ``up_fc_weight``：Tensor类型，必选。 ``fused_norm_residual_ffn`` 算子中升维矩阵权重参数，维度 ``[ffn_inner_size, hidden_size]`` 。
- ``up_fc_bias``：Tensor类型，可选。 ``fused_norm_residual_ffn`` 算子中升维矩阵偏置参数，维度 ``[ffn_inner_size]`` 。其中ffn_inner_size表示FFN中间层的维度。
- ``down_proj_weight``：Tensor类型，必选。 ``fused_norm_residual_ffn`` 算子中降维矩阵权重参数，维度 ``[hidden_size, ffn_inner_size]`` 。
- ``down_proj_bias``：Tensor类型，可选。 ``fused_norm_residual_ffn`` 算子中降维矩阵偏置参数，维度 ``[hidden_size]`` 。
- ``gate_up_proj_weight``：Tensor类型，可选。Gated FFN算子中升维矩阵权重参数，维度 ``[hidden_size, ffn_inner_size]`` 。
- ``gate_up_proj_bias``：Tensor类型，可选。Gated FFN算子中升维矩阵偏置参数，维度 ``[hidden_size]`` 。
- ``norm_weight``：Tensor类型，可选。 ``fused_norm_residual_ffn`` 的layernorm权重参数。维度 ``[input_size]``。
- ``norm_bias``：Tensor类型，可选。 ``fused_norm_residual_ffn`` 的layernorm偏置参数。维度 ``[input_size]``。
- ``eps`` ： float类型。layernorm的 ``eps`` 参数。
- ``act_mode``：string类型，必选。FFN的激化函数类型，可选的激活函数："relu"、"gelu"、"silu"、"none"。
- ``residual_is`` : string类型， 决定 ``residual`` 的来源，当 ``residual_is = input`` 时，表示来自 ``input`` 。
- ``alpha`` ： float类型。默认值为1.0。当 ``residual`` 存在时，作FNN模块的系数。
- ``beta`` ： float类型。默认值为1.0。当 ``residual`` 存在时，作为 ``residual`` 的系数。

**返回值说明**

- output: Tensor类型，必选。 ``ffn`` 的输出。维度 ``[batch_size, seq_lens, hidden_size]`` 。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - up_fc_weight：FP16、BF16
  - up_fc_bias：FP16、BF16
  - down_proj_weight：FP16、BF16
  - down_proj_bias: FP16、BF16
  - gate_up_proj_weight: FP16、BF16
  - gate_up_proj_bias: FP16、BF16
  - norm_weight: FP16、BF16
  - norm_bias: FP16、BF16
  - output: FP16、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_fused_ffn.py`` 。


.. _torch_mlu_ops.fused_rms_norm:

torch_mlu_ops.fused_rms_norm
------------------------------------------

.. code::

  torch_mlu_ops.fused_rms_norm(input,residual,gamma,beta,bias,eps,store_output_before_norm,quant_scale,out)

**功能描述**

RMSnorm归一化算子。

**公式说明**

* RMSnorm标准结构公式如下：

  .. math::

    \begin{aligned}
    output = weight \times \frac{input}{\sqrt{RMS(input) + eps}} + bias\\
    output = Quantize(output, quant_scale)
    \end{aligned}

* 公式中各参数描述如下：

  - ``input``：RMSnorm的输入。
  - ``RMS()``：表示对input取均方（mean square）。
  - ``Quantize()`` 表示对输出做量化。
  - ``weight``：表示RMSnorm的 ``weight`` 参数。
  - ``bias``：表示RMSnorm的 ``bias`` 参数。
  - ``eps``：表示RMSnorm eps的值，一个极小的正数，避免除零，一般默认值设定为0.00001。
  - ``output``：RMSnorm的输出。

**参数说明**

- ``input``：Tensor类型，必选。RMSnorm算子的输入。维度 ``[batch，seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``residual``：Tensor类型，可选。RMSnorm算子的残差输入。在pad模式下，维度 ``[batch，seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``gamma``：Tensor类型，可选。RMSnorm算子的 ``weight`` 。维度 ``[hidden_size]``。
- ``beta``：Tensor类型，可选。RMSnorm算子的 ``bias`` 。维度 ``[hidden_size]``。
- ``bias``：Tensor类型，可选。RMSnorm算子的残差的偏置输入。维度 ``[hidden_size]``。
- ``eps``：Float类型，必选。RMSnorm算子的 ``eps`` 的值，一般设定为0.00001。
- ``store_output_before_norm``：Bool类型，必选。RMSnorm算子是否输出 ``input+residual`` 的值。当为true时，第二个返回值输出 ``input+residual``，当为false时，无第二个返回值。
- ``quant_scale``：Float类型，可选。量化参数。维度 ``[hidden_size]``。当 ``dynamic_quant`` 为False时表示离线量化参数，为True时表示在线量化时的平滑因子。当该参数不为空时，输入输出必须连续。
- ``out``：Tensor类型，可选。RMSnorm算子的原位输出。维度同input，支持batch维度和seq_lens维度的stride。
- ``dynamic_quant``：Bool类型，为False且 ``quant_scale`` 不为空时表示离线per-channel量化输出，否则为在线动态per-token量化输出。

**返回值说明**

- ``output``：Tensor类型，可选。RMSnorm算子的输出，当输入参数out不为空时无此输出。维度 ``[batch, seq_lens, hidden_size]``，支持batch维度和seq_lens维度的stride。
- ``residual_output``：Tensor类型，可选。RMSnorm算子的residual输出。当store_output_before_norm为true时，返回值输出 ``input+residual``，当store_output_before_norm为false时，无此返回值。在pad模式下，维度 ``[batch_size*seq_lens, hidden_size]``；在pack模式下，维度 ``[total_seqlen, hidden_size]``。
- ``smooth_quant_scale``：Tensor类型，可选。维度 ``[batch_size*seq_lens]`` 。当 ``dynamic_quant`` 为True时输出的量化scale。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - residual：FP16、BF16
  - beta：FP16、BF16
  - gamma：FP16、BF16
  - bias：FP16、BF16
  - eps：Float32
  - store_output_before_norm：Bool
  - quant_scale：Float
  - output：FP16、BF16、int8
  - residual_output： FP16、BF16
  - smooth_quant_scale： FP32


**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_fusenorm.py`` 。

.. _torch_mlu_ops.fused_rope:

torch_mlu_ops.fused_rope
------------------------------------------

.. code::

  torch_mlu_ops.fused_rope(input, key_cache_hp, value_cache_hp, sin_cache, cos_cache, position_id, gamma, beta, key_cache_lp, value_cache_lp, cache_bs_id_hp, cache_seq_offsets_hp, cache_bs_id_lp, cache_seq_offsets_lp, k_scale_hp, v_scale_hp, k_scale_lp, v_scale_lp, slot_mapping_hp, slot_mapping_lp, eps)-> None

**功能描述**

fold rope融合算子。

**公式说明**

* 公式如下：

该算子为apply rotary，layernorm，以及reshape linear/paged cache功能的融合。

**参数说明**

- ``qkv``：Tensor类型，必选。输入维度 ``[batch, 1, head_num_q + head_num_k * 2, head_size]`` 。
- ``key_cache_hp``：Tensor类型，必选。当 ``paged=False``时，维度为 ``(max_bs, head_num_k, max_decode_len, head_size)`` 否则 ``(num_blocks, head_num_k, block_size, head_size)`` 。
- ``value_cache_hp``：Tensor类型，必选。当 ``paged=False``时，维度为 ``(max_bs, head_num_k, max_decode_len, head_size)`` 否则 ``(num_blocks, head_num_k, block_size, head_size)`` 。
- ``sin_table``：Tensor类型，必选。输入维度 ``[rotary_seq, head_size]`` 。
- ``cos_table``：Tensor类型，必选。输入维度 ``[rotary_seq, head_size]`` 。
- ``position_id``：Tensor类型，必选。输入维度为 ``[batch]`` 。
- ``gamma``：Tensor类型，必选。输入维度是 ``[head_size]`` 。
- ``beta``：Tensor类型，必选。输入维度是 ``[head_size]`` 。
- ``key_cache_lp``：Tensor类型，可选。当 ``paged=False``时，维度为 ``(max_bs, head_num_k, max_decode_len, head_size / 2)`` 否则 ``(num_blocks, head_num_k, block_size, head_size / 2)`` 。
- ``value_cache_lp``：Tensor类型，可选。当 ``paged=False``时，维度为 ``(max_bs, head_num_k, max_decode_len / 2, head_size)`` 否则 ``(num_blocks, head_num_k, block_size / 2, head_size)`` 。
- ``cache_bs_id_hp``：Tensor类型，可选，默认值为None。当前batch在cache中的位置，维度 ``[batch]`` 。
- ``cache_seq_offsets_hp``：Tensor类型，可选，默认值为None。当前kv在cache中的位置，维度 ``[batch]`` 。
- ``cache_bs_id_lp``：Tensor类型，可选，默认值为None。当前batch在cache中的位置，维度 ``[batch]`` 。
- ``cache_seq_offsets_lp``：Tensor类型，可选，默认值为None。当前kv在cache中的位置，维度 ``[batch]`` 。
- ``key_scale_hp``：Tensor类型，可选，默认值为None。key的量化参数，维度 ``[head_num_k, head_size]`` 。
- ``value_scale_hp``：Tensor类型，可选，默认值为None。value的量化参数，维度 ``[head_num_k, head_size]`` 。
- ``key_scale_lp``：Tensor类型，可选，默认值为None。key的量化参数，维度 ``[max_bs, head_num_k, group_num]`` 。
- ``value_scale_lp``：Tensor类型，可选，默认值为None。value的量化参数，维度 ``[max_bs, head_num_k, group_num]`` 。
- ``slot_mapping_hp``：Tensor类型，可选，默认值为None。paged_cahe时指示当前kv在cache中的位置，维度 ``[batch]`` 。
- ``slot_mapping_lp``：Tensor类型，可选，默认值为None。paged_cahe时指示当前kv在cache中的位置，维度 ``[batch]`` 。
- ``eps``：float32类型，可选。layernorm的eps参数，默认值为1e-5 。

**返回值说明**

-  ``qkv``：原位操作。
-  ``key_cache_hp``：原位操作。
-  ``value_cache_hp``：原位操作。
-  ``key_cache_lp``：原位操作。
-  ``value_cache_lp``：原位操作。
-  ``key_scale_lp``：原位操作。
-  ``value_scale_lp``：原位操作。

**规格限制**

- 数据类型支持：

  - qkv：FP16、BF16
  - key_cache_hp：FP16、BF16、INT8
  - value_cache_hp：FP16、BF16、INT8
  - sin_table：FP16、BF16
  - cos_table：FP16、BF16
  - position_id：INT32
  - gamma：FP16、BF16
  - beta：FP16、BF16
  - key_cache_lp: INT64
  - value_cache_lp: INT4
  - cache_bs_id_hp：INT32
  - cache_seq_offsets_hp：INT32
  - cache_bs_id_lp：INT32
  - cache_seq_offsets_lp：INT32
  - key_scale_hp：FP32
  - value_scale_hp：FP32
  - key_scale_lp：FP32
  - value_scale_lp：FP32
  - slot_mapping_hp：INT32
  - slot_mapping_lp: INT32
  - eps：FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_fused_rope.py`` 。

.. _torch_mlu_ops.group_gemm:

torch_mlu_ops.group_gemm
------------------------------------------

.. code::

  torch_mlu_ops.group_gemm(a, b, m_list, expand_idx, c, alpha, beta, max_m, bias) -> torch.Tensor

**功能描述**

多组GEMM集成算子。

**公式说明**

无

**参数说明**

- ``a``：Tensor类型，必选。GEMM算子的输入矩阵。维度 ``[m, k]``。 ``expand_idx `` 存在时， ``total_m`` 等于 ``m_list`` 元素之和，否则 ``total_m`` 等于 ``m``。
- ``b``：Tensor类型，必选。GEMM算子的权重矩阵。维度 ``[experts_num, n, k]``。
- ``m_list``：Tensor类型，必选。用于指定每个子矩阵的m维度。维度 ``[experts_num]``。
- ``expand_idx``：Tensor类型，可选。用于扩展输入矩阵。维度 ``[total_m]``。
- ``c``：Tensor类型，可选。残差。维度 ``[total_m, n]``。
- ``alpha``：Tensor类型，可选。a乘b的系数。维度 ``[experts_num]``。
- ``beta``：Tensor类型，可选。c的系数。维度 ``[experts_num]``。
- ``max_m``: int类型，必选。m_list中的最大值。由于m_list是device数据，获取真实的最大值会影响性能，建议提供m_list中的理论最大值。
- ``bias``: Tensor类型，可选。GEMM算子的bias矩阵。维度 ``[experts_num, n]``。


**返回值说明**

- ``output``: Tensor类型。 维度 ``[total_m, n]``。

**规格限制**

- 数据类型支持：

  - a：HALF、FP32、BF16、INT8
  - b：HALF、FP32、BF16、INT8
  - m_list: INT32
  - expand_idx: INT32
  - c：HALF、FP32、BF16、INT8
  - alhpa：FP32
  - beta：FP32
  - output：HALF、FP32、BF16
  - bias: HALF、FP32、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_group_gemm.py`` 。

.. _torch_mlu_ops.smooth_quant_group_gemm:

torch_mlu_ops.smooth_quant_group_gemm
------------------------------------------

.. code::

  torch_mlu_ops.smooth_quant_group_gemm(a, b, m_list, expand_idx, c, alpha, beta, a_scale, b_scale, dtype, max_m, bias, quant_flag) -> torch.Tensor

**功能描述**

多组GEMM集成量化功能算子。

**公式说明**

无

**参数说明**

- ``a``：Tensor类型，必选。GEMM算子的输入矩阵。维度 ``[m, k]``。 ``expand_idx `` 存在时， ``total_m`` 等于 ``m_list`` 元素之和，否则 ``total_m`` 等于 ``m``。
- ``b``：Tensor类型，必选。GEMM算子的权重矩阵。维度 ``[experts_num, n, k]``。
- ``m_list``：Tensor类型，必选。用于指定每个子矩阵的m维度。维度 ``[experts_num]``。
- ``expand_idx``：Tensor类型，可选。用于扩展输入矩阵。维度 ``[total_m]``。
- ``c``：Tensor类型，可选。残差。维度 ``[total_m, n]``。
- ``alpha``：Tensor类型，可选。a乘b的系数。维度 ``[experts_num]``。
- ``beta``：Tensor类型，可选。c的系数。维度 ``[experts_num]``。
- ``a_scale``：Tensor类型，可选。输入的分token量化的scale系数。维度 ``[total_m]``。
- ``b_scale``：Tensor类型，可选。权重的分channel量化的scale系数。维度 ``[experts_num, n]`` 非group量化场景，或 ``[quant_group, experts_num, n]`` group量化场景。
- ``dtype``：dype类型， ``output`` 的数据类型，必须是 ``torch.half`` 、 ``torch.bfloat16`` 。
- ``max_m``: int类型，必选。m_list中的最大值。由于m_list是device数据，获取真实的最大值会影响性能，建议提供m_list中的理论最大值。
- ``bias``: Tensor类型，可选。GEMM算子的bias矩阵。维度 ``[experts_num, n]``。
- ``quant_flag``：List类型，可选。w4w8混合量化模式时标记quant_group的bit位，数值必须是4或8，其它值会引发错误。长度 ``[experts_num * quant_group]``。

**返回值说明**

- ``output``: Tensor类型。 维度 ``[total_m, n]``。

**规格限制**

- 数据类型支持：

  - a：HALF、FP32、BF16、INT8
  - b：HALF、FP32、BF16、INT8、INT8(INT4x2)
  - m_list: INT32
  - expand_idx: INT32
  - c：HALF、FP32、BF16、INT8
  - alhpa：FP32
  - beta：FP32
  - a_scale: FP32
  - b_scale: FP32
  - output：HALF、FP32、BF16
  - bias: HALF、FP32、BF16

**备注**

- group量化场景，``quant_group * quant_group_size == k``，``quant_group_size`` 表示该场景下每个group的元素数量，``quant_group_size`` 可以取值128、256、512、1024。
- group量化场景， b_scale的形状必须是(quant_group, experts_num, n)。
- gather 模式不支持int4量化。
- 4D 转数模式不支持int4量化。
- 不支持通用量化模式的int4量化。
- ``quant_flag`` 存在时为w4w8混合量化模式，其元素只能是4或8。4表示该quant_group对应的权重是int4量化，8表示该quant_group对应的权重是int8量化。
- w4w8混合量化模式要求b必须是1D。

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_group_gemm.py`` 。

.. _torch_mlu_ops.matmul:

torch_mlu_ops.matmul
------------------------------------------

  torch_mlu_ops.matmul(input, filter, bias, residual, act_mode, alpha, beta, use_fast, approximate, out_dtype, a_scale, b_scale, trans_a, trans_b)-> torch.Tensor

**功能描述**

通用的矩阵乘法算子。

**公式说明**

* 带后融合的matmul标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     mul_out = alpha * (matmul(input, filter) + bias) + beta * residual\\
     output = active(mul_out)\\
     \end{array}

* 公式中各参数描述如下：

  - ``input`` ：matmul的左矩阵。
  - ``filter``：matmul的右矩阵。
  - ``bias``：偏置。
  - ``residual``：参差。
  - ``output`` ：计算最终的输出。

* 公式中各函数描述如下：

  - ``matmul`` ：矩阵乘。
  - ``active``：激活计算。

**参数说明**

- ``input``：Tensor类型，必选。矩阵乘的左矩阵。维度 ``[mat_m, mat_k]``。
- ``filter``：Tensor类型，必选。矩阵乘的右矩阵。维度 ``[mat_n, mat_k]``。
- ``bias``：Tensor类型，可选。矩阵乘的偏置。维度 ``[1, mat_n]``。
- ``residual``：Tensor类型，可选。参差。维度 ``[mat_m, mat_n]``。
- ``act_mode``：String类型，可选。后融合激活的类型，可选的激活函数： ``relu`` 、 ``gelu`` 、 ``silu``。
- ``alpha``：Float32类型，必选。矩阵乘法系数。
- ``beta``：Float32类型，必选。参差的系数。
- ``use_fast``：Bool类型，必选。是否使用高性能的方式计算激活。
- ``approximate``：Bool类型，必选。gelu是否使用更精确的近似计算方式。
- ``out_dtype``：String类型，可选。如果输入为INT8类型，则必须指定输出类型。
- ``a_scale``：Float32类型，默认值1.0。当input为INT8类型时对应的量化系数。
- ``b_scale``：Float32类型，默认值1.0。当weight为INT8类型时对应的量化系数。
- ``trans_a``：Bool类型，默认值为False，标记input是否需要转置。
- ``trans_b``：Bool类型，默认值为True，标记weight是否需要转置。


**返回值说明**

- ``out``：Tensor类型，算子的输出，维度 ``[mat_m, mat_n]``。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16、FP32、INT8
  - filter：FP16、BF16、FP32、INT8
  - bias：FP16、BF16、FP32
  - residual：FP16、BF16、FP32
  - out：FP16、BF16、FP32
  - alpha：FP32
  - beta：FP32
  - use_fast：BOOL
  - approximate：BOOL
  - a_scale：FP32
  - b_scale：FP32
  - trans_a：BOOL
  - trans_b：BOOL

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_matmul.py`` 。

.. _torch_mlu_ops.matmul_allreduce:

torch_mlu_ops.matmul_allreduce
------------------------------------------

.. code::

  torch_mlu_ops.matmul_allreduce(cncl_comm, input, filter, bias, residual, alpha, beta, block_m)-> torch.Tensor

**功能描述**

矩阵乘法与allreduce的通信计算融合算子。

**公式说明**

无

**参数说明**

- ``cncl_comm``：Int类型，必选。cnclComm，用于归约通信。
- ``input``：Tensor类型，必选。矩阵乘的左矩阵。维度 ``[mat_m, mat_k]``。
- ``filter``：Tensor类型，必选。矩阵乘的右矩阵。维度 ``[mat_n, mat_k]``。
- ``bias``：Tensor类型，可选。矩阵乘的偏置。维度 ``[1, mat_n]``。
- ``residual``：Tensor类型，可选。参差。维度 ``[mat_m, mat_n]``。
- ``alpha``：Float32类型，必选。矩阵乘法系数。
- ``beta``：Float32类型，必选。参差的系数。
- ``block_m``：Int类型，可选。input在mat_m维度拆分的份数。

**返回值说明**

- ``out``：Tensor类型，算子的输出，维度 ``[mat_m, mat_n]``。

**规格限制**

- 数据类型支持：

  - cncl_comm: INT64
  - input：FP16、BF16、FP32
  - filter：FP16、BF16、FP32
  - bias：FP16、BF16、FP32
  - residual：FP16、BF16、FP32
  - out：FP16、BF16、FP32
  - alpha：FP32
  - beta：FP32
  - block_m: INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/allreduce_case/test_matmul_all_reduce.py`` 。

.. _torch_mlu_ops.moe_combine_result:

torch_mlu_ops.moe_combine_result
-------------------------------------

.. code::

  torch_mlu_ops.moe_combine_result(input, reduce_weight, gather_ids, residual(optional),
                                   cusum_token_count(optional), start_expert_id, expert_size,
                                   bias(optional))->Tensor


**功能描述**

融合算子。从不同的专家模型中收集计算结果 (gather) ，然后对这些结果进行加权组合 (reduce) 。

**公式说明**

无。

**参数说明**

- ``input``：Tensor类型，必选。算子的输入，维度 ``[num_tokens, hidden_size]`` 。其中 ``num_tokens = num_token * topk`` ，``num_token`` 表示token数量，``topk`` 表示每个token选取概率在前topk的专家，``hidden_size`` 表示隐向量维数。
- ``reduce_weight``：Tensor类型，必选。token的权重信息，维度 ``[num_tokens]`` 。
- ``gather_ids``：Tensor类型，必选。用于收集token的索引，维度 ``[num_tokens]`` 。
- ``residual``：Tensor类型，可选，默认值为None。表示残差，维度 ``[num_token, hidden_size]`` 。
- ``cusum_token_count``：Tensor类型，可选，默认值为None。token_count的前缀和，维度 ``[num_expert + 1]`` ，其中 ``num_expert`` 表示专家数量。
- ``start_expert_id``：Int32类型，必选。表示EP模型下起始专家的位置信息。
- ``expert_size``：Int32类型，必选。表示EP模型下实际使用的专家数量。
- ``bias``：Tensor类型，可选，默认值为。表示偏置，维度 ``[num_expert, hidden_size]`` 。

**返回值说明**

- output: Tensor类型，必选。维度为 ``[num_token, hidden_size]`` 。

**规格限制**

- 数据类型支持：

  - input：FP16、FP32、BF16
  - reduce_weight：FP32
  - gather_ids：INT32
  - residual：FP16、FP32、BF16
  - cusum_token_count：INT32
  - start_expert_id：INT32
  - expert_size：INT32
  - bias：FP16、FP32、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_combine_result.py`` 。

.. _torch_mlu_ops.moe_expand_input:

torch_mlu_ops.moe_expand_input
------------------------------------------

.. code::

  torch_mlu_ops.moe_expand_input(input, gather_idx, cusum_token_count, output, start_expert_id, expert_size)->Tensor

**功能描述**

根据 ``gather_idx`` 在 ``input`` 执行索引操作。其中 ``cusum_token_count`` 、 ``start_expert_id``、 ``expert_size`` 在EP模式下使用，用于确定当前专家处理token的index范围。

**公式说明**

无。

**参数说明**

- ``input``：Tensor类型，必选。算子的输入，输入维度 ``[token_num, hidden_size]``。其中 ``token_num`` 表示token数量， ``hidden_size`` 表示各个token的特征信息。
- ``gather_idx``：Tensor类型，必选。输入维度 ``[token_num * topk]``， ``topk`` 表示每个token选取概率在前topk的专家进行计算。
- ``cusum_token_count``：Tensor类型，可选，默认值为None。输入维度为 ``[expert_num + 1]`` ，其中 ``expert_num`` 表示专家的数量。
- ``output``：Tensor类型，可选，默认值为None。算子的原位输出，维度为 ``[token_num * topk, hidden_size]`` 。
- ``start_expert_id``：Int32类型，可选，默认值为0。表示EP模型下起始专家的位置信息。
- ``expert_size``：Int32类型，可选，默认值为0。表示EP模式下实际使用的专家数量。

**返回值说明**

- output，可选。维度为 ``[token_num * topk, hidden_size]``。

**规格限制**

- 数据类型支持：

  - input：INT8，FP16，FP32，BF16
  - gather_idx：INT32
  - cusum_token_count：INT32
  - output：INT8，FP16，FP32，BF16
  - start_expert_id：INT32
  - expert_size：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_expand_input.py`` 。

.. _torch_mlu_ops.moe_gen_idx:

torch_mlu_ops.moe_gen_idx
------------------------------------------

.. code::

  torch_mlu_ops.moe_gen_idx(expert_id, expert_num)->Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

**功能描述**

基于 ``expert_id`` 和 ``expert_num`` 生成 ``expand_idx`` 、 ``combine_idx`` 、 ``token_count`` 、 ``cusum_token_count`` 。

**公式说明**

无。

**参数说明**

- ``expert_id``：Tensor类型，必选。算子的输入，输入维度 ``[token_num, topk]``。其中 ``token_num`` 表示token数量， ``topk`` 表示每个token选取概率在前topk的专家进行计算。
- ``expert_num``：Int32类型，必选。算子的输入，表示专家的数量。

**返回值说明**

- outputs，Tuple[Tensor, Tensor, Tensor, Tensor]类型，可选。输出的4个Tensor分别是expand_idx、combine_idx、token_count和cusum_token_count，维度同上。

**规格限制**

- 数据类型支持：

  - expert_id：INT32
  - expert_num：INT32
  - outputs：输出的4个Tensor类型均为INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_gen_idx.py`` 。

.. _torch_mlu_ops.moe_softmax_topk:

torch_mlu_ops.moe_softmax_topk
------------------------------------------

.. code::

  torch_mlu_ops.moe_softmax_topk(input,topk,num_expert_group,topk_group,normalize,mask,normed_by)->Tuple[torch.Tensor]

**功能描述**

MOE Route模块中的softmax topk融合算子。

**公式说明**

无。

**参数说明**

- ``input``：Tensor类型，必选。算子的输入，输入维度 ``[..., expert_num]``。需要保证连续性。
- ``topk``：INT32类型，每个token选取的专家数。
- ``num_expert_group``：INT32类型，表示每个专家的分组个数。
- ``topk_group``：INT32类型，每个专家组的topk个数。
- ``normalize``：Bool类型，是否给 ``reduce_weight`` 添加归一化操作。
- ``mask``：Tensor类型，可选。用于与softmax的计算结果相乘, 维度数量与input一致，后两维度大小与input一致。需要保证连续性。
- ``normed_by``：Str类型，仅在normalize为True的时候使能，用于表示归一化的模式(topk_logit或softmax_logit)。

**返回值说明**

- ``outputs`` : 可选，Tuple[Tensor, Tensor]类型。输出的两个Tensor分别是 ``reduce_weight`` 和 ``expert_id`` 。

**规格限制**

- 数据类型支持：

  - input ：Float，half，bfloat16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_softmax_topk.py`` 。


.. _torch_mlu_ops.preload:

torch_mlu_ops.preload
------------------------------------------

.. code::

  torch_mlu_ops.preload(weight,size)

**功能描述**

预加载权重。

**公式说明**

无。

**参数说明**

- ``weight``：Tensor类型，必选。预加载的权重。
- ``size``：int类型，必选。预加载的数据大小（字节）。

**返回值说明**

无返回值。

**规格限制**

无。

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_preload.py`` 。


.. _torch_mlu_ops.quant_to_linear_cache:

torch_mlu_ops.quant_to_linear_cache
------------------------------------------

.. code::

  torch_mlu_ops.quant_to_linear_cache(key,value,key_cache,value_cache,key_cache_scale,value_cache_scale,context_lengths,max_context_len,packed,context_seq_offset,cache_bs_id,cache_seqlen_offset,quant_bit)

**功能描述**

将计算得到的K-V拼接到K-V cache上的INT8量化版本算子。


**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[batch_size, max_context_len, head_num, head_size]``。
- ``value``：Tensor类型，可选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[batch_size, max_context_len, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[max_batch, head_num, cache_mem_len, head_size]``。其中max_batch要求大于batch_size，max_context_len为网络支持的最大sequences长度。
- ``value_cache``：Tensor类型，可选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为 ``[max_batch, head_num, cache_mem_len, head_size]``。
- ``key_cache_scale``：Tensor类型，必选。Transformer模块中保存的 ``Key`` 的 ``scale`` 值。在context阶段，维度为 ``[batch_size, head_num, cache_mem_len]`` 或者 ``[batch_size, head_num, cache_mem_len, group_num]``。
- ``value_cache_scale``：Tensor类型，可选。Transformer模块中保存的 ``Value`` 的 ``scale`` 值。在context阶段，维度为 ``[batch_size, head_num, cache_mem_len]`` 或者 ``[batch_size, head_num, cache_mem_len, group_num]``。
- ``context_lengths``：Tensor类型，必选。为每批batch的sequences length长度。当 ``packed=false`` 时，维度为 ``[batch_size]``，每个batch里为对于的sequence length；当 ``packed=true`` 时，维度为 ``[batch_size+1]``，用sequence与sequence之间的长度差表示sequence长度。
- ``max_context_len``：Int类型，必选。在context阶段，表示最大的sequence length。
- ``packed``：Bool类型，必选。当 ``packed=true`` 时，表示开启pack模式，当 ``packed=false`` 时，表示关闭packed模式。
- ``context_seq_offset``：Tensor类型，可选。表示context的seq的偏置，维度 ``[batch_size]``。
- ``cache_bs_id``：Tensor类型，可选。表示cache的batch的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 1, 2 ... batch - 1}。
- ``cache_seqlen_offset``：Tensor类型，可选。表示cache的seq_lens的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 0, 0 ... 0}。
- ``quant_bit``：Int类型，可选。表示量化结果的位宽，可选8（表示int8）或4（表示int4），默认值为8。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16
  - value：FP16、BF16
  - key_cache：INT8
  - value_cache：INT8
  - key_cache_scale： FP32
  - value_cache_scale： FP32
  - context_lengths：INT32
  - max_context_len：INT32
  - packed：Bool
  - context_seq_offset：INT32
  - cache_bs_id： INT32
  - cache_seqlen_offset： INT32
  - quant_bit : INT32
  - group_size: INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_quant_to_linear_cache.py`` 。

.. _torch_mlu_ops.quant_to_paged_cache:

torch_mlu_ops.quant_to_paged_cache
------------------------------------------

.. code::

  torch_mlu_ops.quant_to_paged_cache(key,value,key_cache,value_cache,key_cache_scale,value_cache_scale,slot_mapping)

**功能描述**

:ref:`torch_mlu_ops.reshape_paged_cache` 的INT8量化形式算子。

**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``value``：Tensor类型，必选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。
- ``value_cache``：Tensor类型，必选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。
- ``key_cache_scale``：Tensor类型，必选。Transformer模块中保存 ``Key cache scale`` 的Tensor。维度为 ``[num_blocks, head_num, block_size]``。
- ``value_cache_scale``：Tensor类型，必选。Transformer模块中保存 ``Value cache scale`` 的Tensor。维度为 ``[num_blocks, head_num, block_size]``。
- ``slot_mapping``：Tensor类型，必选。表示Paged Attention中 ``K-V cache`` 存放位置。维度为 ``[num_tokens]``。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16, FP32
  - value：FP16、BF16, FP32
  - key_cache：INT8
  - value_cache：INT8
  - key_cache_scale：FP32
  - value_cache_scale：FP32
  - slot_mapping：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/kernel_pytest/quant_paged_cache_kernel_test/quant_paged_cache_kernel_test.py``。

.. _torch_mlu_ops.offline_quant_to_linear_cache:

torch_mlu_ops.offline_quant_to_linear_cache
-------------------------------------------------

.. code::

  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)

**功能描述**

:ref:`torch_mlu_ops.quant_to_linear_cache` 的离线量化形式算子。

**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。 如果使用pack模式， 维度为 ``[num_tokens, head_num_kv, head_size]``， 如果使用pad模式， 维度为 ``[batch, seqlen, head_num_kv, head_size]``。
- ``value``：Tensor类型，可选。Transformer模块中计算得到的 ``Value``。 维度与key相同。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[max_batch, head_num_kv, cache_mem_len, head_size]``。
- ``value_cache``：Tensor类型，可选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为与key_cache相同。
- ``key_cache_quant_scale``：Tensor类型，必选。Transformer模块中保存 ``Key cache scale`` 的Tensor。当量化模式是per_channel时，维度为 ``[head_num_kv, head_size]``, 当量化模式是per_head时，维度是 ``[head_num_kv, cache_mem_len]``。
- ``value_cache_quant_scale``：Tensor类型，可选。Transformer模块中保存 ``Value cache scale`` 的Tensor。维度与value_cache_quant_scale相同。
- ``context_lengths``：Tensor类型，必选。表示context的长度。如果使用pack模式，维度为 ``[batch + 1]``，如果使用pad模式，维度为 ``[batch]``。
- ``max_context_len``：Int类型，必选。表示context的最大长度。
- ``quant_mode``：int类型，必选。表示量化模式。当该参数为True时，表示使用per_head量化，否则使用per_channel量化。
- ``packed``：bool类型，必选。该值为True时，表示输入为pack模式，否则为pad模式。
- ``context_seq_offset``：Tensor类型，可选。表示待量化的key或value的位置索引。维度为 ``[batch]``。
- ``cache_bs_id``：Tensor类型，可选。表示key或value放入的cache的batch索引。维度为 ``[batch]``。
- ``cache_seqlen_offset``：Tensor类型，可选。表示key或value放入cache的起始位置。维度为 ``[batch]``。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16, FP32
  - value：FP16、BF16, FP32
  - key_cache：INT8
  - value_cache：INT8
  - key_cache_quant_scale：FP32
  - value_cache_quant_scale：FP32
  - context_lengths：INT32
  - max_context_len：INT32
  - quant_mode：INT32
  - packed：BOOL
  - context_seq_offset：INT32
  - cache_bs_id：INT32
  - cache_seqlen_offset：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/kernel_pytest/test_offline_quant_to_linear_cache.py``。

.. _torch_mlu_ops.offline_quant_to_paged_cache:

torch_mlu_ops.offline_quant_to_paged_cache
------------------------------------------

.. code::

  torch_mlu_ops.offline_quant_to_paged_cache(key,value,key_cache_scale,value_cache_scale,slot_mapping,key_cache,value_cache)

**功能描述**

:ref:`torch_mlu_ops.quant_to_paged_cache` 的离线量化形式算子。

**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``value``：Tensor类型，必选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``key_cache_scale``：Tensor类型，必选。Transformer模块中保存 ``Key cache scale`` 的Tensor。维度为 ``[head_num, head_size]``。
- ``value_cache_scale``：Tensor类型，必选。Transformer模块中保存 ``Value cache scale`` 的Tensor。维度为 ``[head_num, head_size]``。
- ``slot_mapping``：Tensor类型，必选。表示Paged Attention中 ``K-V cache`` 存放位置。维度为 ``[num_tokens]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。
- ``value_cache``：Tensor类型，必选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16, FP32
  - value：FP16、BF16, FP32
  - key_cache_scale：FP32
  - value_cache_scale：FP32
  - slot_mapping：INT32
  - key_cache：INT8
  - value_cache：INT8

**使用示例**

请参考 ``torch_mlu_ops/tests/kernel_pytest/test_offline_quant_to_paged_cache_kernel.py``。

.. _torch_mlu_ops.reshape_linear_cache:

torch_mlu_ops.reshape_linear_cache
------------------------------------------

.. code::

  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)

**功能描述**

将计算得到的K-V拼接到K-V cache上的算子。

**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[batch_size, max_context_len, head_num, head_size]``。
- ``value``：Tensor类型，可选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[batch_size, max_context_len, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[max_batch, head_num, cache_mem_len, head_size]``。其中 ``max_batch`` 要求大于 ``batch_size``，``max_context_len`` 为网络支持的最大sequences长度。
- ``value_cache``：Tensor类型，可选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为 ``[max_batch, head_num, cache_mem_len, head_size]``。
- ``context_lengths``：Tensor类型，必选。为每批batch的sequences length长度。当 ``packed=false`` 时，维度为 ``[batch_size]``，每个batch里为对于的sequence length；当 ``packed=true`` 时，维度为 ``[batch_size+1]``，用sequence与sequence之间的长度差表示sequence长度。
- ``max_context_len``：Int类型，必选。在context阶段，表示最大的sequence length。
- ``packed``：Bool类型，必选。当 ``packed=true`` 时，表示开启pack模式，当 ``packed=false`` 时，表示关闭pack模式。
- ``context_seq_offset``：Tensor类型，可选。表示context的seq的偏置，维度 ``[batch_size]``。
- ``cache_bs_id``：Tensor类型，可选。表示cache的batch的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 1, 2 ... batch - 1}。
- ``cache_seqlen_offset``：Tensor类型，可选。表示cache的seq_lens的偏置，维度 ``[batch_size]``。如果为空，则默认值为{0, 0, 0 ... 0}。


**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16
  - value：FP16、BF16
  - key_cache：FP16、BF16
  - value_cache：FP16、BF16
  - context_lengths：INT32
  - max_context_len：INT32
  - packed：Bool
  - context_seq_offset：INT32
  - cache_bs_id： INT32
  - cache_seqlen_offset： INT32


**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_reshape_linear_cache.py`` 。


.. _torch_mlu_ops.reshape_paged_cache:

torch_mlu_ops.reshape_paged_cache
------------------------------------------

.. code::

  torch_mlu_ops.reshape_paged_cache(key,value,key_cache,value_cache,slot_mapping)

**功能描述**

:ref:`torch_mlu_ops.reshape_linear_cache` 的Paged形式算子。

**公式说明**

无。

**参数说明**

- ``key``：Tensor类型，必选。Transformer模块中计算得到的 ``Key``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``value``：Tensor类型，必选。Transformer模块中计算得到的 ``Value``。在context阶段，维度为 ``[num_tokens, head_num, head_size]``。
- ``key_cache``：Tensor类型，必选。Transformer模块中保存 ``Key cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。
- ``value_cache``：Tensor类型，必选。Transformer模块中保存 ``Value cache`` 的Tensor。维度为 ``[num_blocks, head_num, block_size, head_size]``。
- ``slot_mapping``：Tensor类型，必选。表示Paged Attention中 ``K-V cache`` 存放位置。维度为 ``[num_tokens]``。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - key：FP16、BF16
  - value：FP16、BF16
  - key_cache：FP16、BF16
  - value_cache：FP16、BF16
  - slot_mapping：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_reshape_paged_cache.py`` 。


.. _torch_mlu_ops.quantize:

torch_mlu_ops.quantize
------------------------------------------

.. code::

  torch_mlu_ops.quantize(x, scale, zero=None)-> torch.Tensor

**功能描述**

对输入tensor进行量化的算子。

**公式说明**

* 公式如下：

  .. math::

      \begin{array}{lcl}
      scaled = x * scale\\
      output = scaled.round().clamp(-128, 127).to(torch.int8)
      \end{array}

* 公式中各参数描述如下：

  - ``x`` ：quantize算子的输入。
  - ``scale``：量化过程的scale参数。
  - ``output``： quantize算子的输出。

**参数说明**

- ``x`` ：Tensor类型，必选。quantize算子的输入，维度是 ``[..., C]``，其中 ``C`` 是通道的维度。
- ``scale`` ：Tensor类型，必选。quantize算子的 ``scale``输入，维度是 ``[C]``。
- ``zero`` ：Tensor类型，必选。quantize算子的保留参数，填入 ``None``。

**返回值说明**

- ``output``：Tensor类型，必选。算子的输出，维度和 ``x`` 一致。

**规格限制**

- 数据类型支持：

  - x：FP32、FP16、BF16
  - scale：FP32
  - output：INT8

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_smooth_quant.py`` 。

.. _torch_mlu_ops.single_query_cached_kv_attn:

torch_mlu_ops.single_query_cached_kv_attn
------------------------------------------

.. code::

  torch_mlu_ops.single_query_cached_kv_attn(q_ori,k_cache,v_cache,out,block_tables,context_lens,k_cache_quant_scale,v_cache_quant_scale,alibi_slopes,max_context_len,windows_size_left,windows_size_right,softmax_scale,return_lse,kv_cache_quant_bit_size)->Tensor

**功能描述**

Decoder阶段的Single Query Cached KV Attention算子。

**公式说明**

* Single Query Cached KV Attention标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     scores = batch\_matmul(q\_ori, k\_cache, transpose\_b=True) * softmax\_scale\\
     scores +=  ALiBi\\
     scores += causal\_mask\\
     attention = softmax(scores)\\
     output = batch\_matmul(attention, v\_cache)\\
     lse = logsumexp(scores)\\
     \end{array}

* 公式中各参数描述如下：

  - ``q_ori`` ：Single Query Cached KV Attention模块的输入。
  - ``ALiBi``：在线计算的ALiBi位置编码矩阵。
  - ``scores`` ：Single Query Cached KV Attention模块第一个批量矩阵乘的输出。
  - ``causal_mask``：因果mask， 当seq_q>1时默认使能。
  - ``attention`` ：Single Query Cached KV Attention模块经过sotmax后的输出。
  - ``output`` ：Single Query Cached KV Attention模块第二个批量矩阵乘的输出，也是模块的最终输出。
  - ``lse`` ：Single Query Cached KV Attention的lse输出。

* 公式中各函数描述如下：

  - ``batch_matmul`` ：批量矩阵乘。
  - ``softmax`` ：归一化指数函数。
  - ``logsumexp`` ：对数求和指数函数。

**参数说明**

- ``q_ori``：Tensor类型，必选。算子的输入。维度 ``[batch, seq_q, q_head_num, head_size_qk]``。
- ``k_cache``：Tensor类型，必选。算子的输入。维度 ``[num_blocks, kv_head_num, block_size, head_size_qk]``。
- ``v_cache``：Tensor类型，必选。算子的输入。维度 ``[num_blocks, kv_head_num, block_size, head_size_v]``。
- ``out``：Tensor类型，可选。维度同input。
- ``block_tables``：Tensor类型，必选。block的id与kv_cache之间的映射关系。维度 ``[bs, max_num_blocks_per_seq]``，非paged attention时，max_num_blocks_per_seq为1。
- ``context_lens``：Tensor类型，必选。记录每个batch的实际seq长度，维度 ``[batch]``。
- ``k_cache_quant_scale``：Tensor类型，可选。k_cache&v_cache为量化类型时k_cache的量化参数，维度 ``[num_blocks, kv_head_num, block_size]``表示per_token量化、 ``[num_blocks, kv_head_num, block_size, 1]`` 表示带group的per_token量化，但当前group_num必须为1， 以及 ``[kv_head_num, head_size_qk]`` 表示per_channel量化。
- ``v_cache_quant_scale``：Tensor类型，可选。k_cache&v_cache为量化类型时v_cache的量化参数，维度同 ``k_cache_quant_scale``。
- ``alibi_slope``：Tensor类型，可选。是一种相对位置编码方法，维度 ``[batch, q_head_num]``，类型必须是float32。
- ``max_context_len``：Int64类型，必选。算子的最大seq长度。
- ``window_size_left``：Int64类型，左侧窗口大小，未指定时设置-1。实际参与计算的key和value长度为 ``window_size_left + 1``。
- ``window_size_right``：Int64类型，保留字段，设置-1即可。
- ``softmax_scale``：Float32类型，必选。softmax之前的系数。
- ``return_lse``：Bool类型，标记是否返回lse tensor。为True时当前仅支持seq_q=1。
- ``kv_cache_quant_bit_size``：INT32类型，KV cache量化的位宽，可选8（表示int8）、4（表示int4）和-1（表示和input一样的数据类型），默认值为-1

**返回值说明**

- ``output``：Tensor类型，可选。算子的输出，当out不为空时有此输出，维度 ``[batch, seq_q, q_head_num, head_size]``。
- ``output_lse``：Tensor类型，可选。算子的lse输出，当 ``return_lse`` 为True时有此输出，维度 ``[batch, q_head_num, seq_q]``。

**规格限制**

- 数据类型支持：

  - q_ori：FP16、BF16、FP32
  - k_cache：FP16、BF16、FP32、INT8、int4
  - v_cache：FP16、BF16、FP32、INT8、int4
  - block_tables：INT32
  - context_lens：INT32
  - k_cache_quant_scale：FP32
  - v_cache_quant_scale：FP32
  - alibi_slope：FP32
  - max_context_len：INT64
  - window_size_left：INT32
  - window_size_right：INT32
  - softmax_scale：FP32
  - output：FP16、BF16、FP32
  - output_lse：FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_single_query_cached_kv_attn.py`` 。

.. _torch_mlu_ops.smooth_quant_matmul:

torch_mlu_ops.smooth_quant_matmul
------------------------------------------

.. code::

  torch_mlu_ops.smooth_quant_matmul(a,a_scale,b,b_scale,dtype,bias,c,act_mode,alpha,beta,use_hp_active)-> torch.Tensor

**功能描述**

实现smooth_quant量化功能的GEMM算子。

**公式说明**

* smooth_quant_matmul标准结构公式如下：

  .. math::

     \begin{array}{lcl}
        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
     \end{array}


* 公式中各参数描述如下：

  - ``a`` ：矩阵乘运算的a矩阵。
  - ``a_scale`` : 输入的分token量化的scale系数
  - ``b`` ：矩阵乘运算的b矩阵。
  - ``b_scale`` : 权重的分channel量化的scale系数
  - ``bias`` ：偏置。
  - ``alpha`` ：浮点标量，默认1.0f。
  - ``beta`` ：浮点标量，默认1.0f。
  - ``c`` ：第三个elemenwise加c矩阵。
  - ``ouput`` :GEMM返回值。

* 公式中各函数描述如下：
  - ``matmul`` :矩阵乘操作。
  - ``active`` :激活操作。

**参数说明**

- ``a``：Tensor类型，必选。GEMM算子的输入。维度 ``[M, K]``。
- ``a_scale``：Tensor类型，可选。输入的分token量化的scale系数。维度 ``[N]``。
- ``b``：Tensor类型，必选。GEMM算子的权重。维度 ``[N, K]``。
- ``b_scale``：Tensor类型，可选。权重的分channel量化的scale系数。维度 ``[N]``。
- ``dtype``：dype类型， ``output`` 的数据类型，必须是 ``torch.float`` 、 ``torch.half`` 、 ``torch.bfloat16`` 。
- ``bias``：Tensor类型，可选。GEMM算子的偏置。维度 ``[N]``。
- ``c`` ： Tensor类型，可选。维度 ``[M, N]``。
- ``act_mode`` : string类型，可选的激活函数："relu"、"gelu"、"silu"、"none"。
- ``alpha`` ： float类型。默认值为1.0。
- ``beta`` ： float类型。默认值为1.0。
- ``use_hp_active`` ：bool类型，激活的计算方式，默认值为False。


**返回值说明**

- ``ouput``: Tensor类型, 维度 ``[M, N]``。

**规格限制**


- 数据类型支持：

  - a：INT8
  - a_scale：FP32
  - b: INT8
  - b_scale：FP32
  - bias：HALF、FP32、BF16
  - c：HALF、FP32、BF16
  - output：HALF、FP32、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_quant_ffn.py`` 。


.. _torch_mlu_ops.smooth_quant_matmul_allreduce:

torch_mlu_ops.smooth_quant_matmul_allreduce
------------------------------------------------

.. code::

  torch_mlu_ops.smooth_quant_matmul_allreduce(cncl_comm, a, a_scale, b, b_scale, dtype, bias, c, alpha, beta, block_m)

**功能描述**

集成平滑量化功能的GEMM通信计算融合算子。

**公式说明**

无。

**参数说明**

- ``cncl_comm``：Int类型，必选。cnclComm，用于归约通信。
- ``a``：Tensor类型，必选。GEMM算子的输入。维度 ``[M, K]``。
- ``a_scale``：Tensor类型，必选。输入的scale系数。维度 ``[M]``。
- ``b``：Tensor类型，必选。GEMM算子的权重。维度 ``[N, K]``。
- ``b_scale``：Tensor类型，必选。权重的scale系数。维度 ``[N]``。
- ``dtype``：Torch.dtype类型，必选。返回值类型，可选的类型： ``torch.half``、 ``torch.bfloat16``, ``torch.float32``。
- ``bias``：Tensor类型，可选。GEMM算子的偏置。维度 ``[N]``。
- ``c``：Tensor类型，可选。残差。维度 ``[M, N]``。
- ``alpha``：float类型，可选。默认1.0f。
- ``beta``：float类型，可选。默认1.0f。
- ``block_m``：Int类型，可选。a在M维拆分的份数。

**返回值说明**

- ``d``: Tensor类型, 维度 ``[M, N]``。

**规格限制**

- 参数限制：

  - M/N/K: M<=INT32_MAX, NK<=INT32_MAX, K>=16, N>=16。
  - a_scale/b_scale: 维度分别是 ``[M]`` 和 ``[N]``。

- 数据类型支持：

  - cncl_comm：INT64
  - a：HALF、FP32、BF16、INT8
  - a_scale：FP32
  - b：INT8
  - b_scale：FP32
  - bias：HALF、FP32、BF16
  - c：HALF、FP32、BF16
  - block_m: INT32
  - d：HALF、FP32、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/allreduce_case/test_quant_matmul_all_reduce.py`` 。

.. _torch_mlu_ops.swap_blocks:

torch_mlu_ops.swap_blocks
------------------------------------------

.. code::

  torch_mlu_ops.swap_blocks(dst, src, block_mapping)

**功能描述**

将 ``src`` 中源位置的BLOCK复制到 ``dst`` 的指定位置，拷贝方向可以是：HostToDev、DevToHost、DevToDev。

**公式说明**

无。

**参数说明**

- ``src``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, block_size, head_size>]``。其中 ``num_layers`` 表示layer数量， ``[num_heads, block_size, head_size]`` 整体大小作为每次需要copy的BLOCK大小。 ``num_blocks`` 表示BLOCK数量； ``block_size`` 表示输入序列中的特定块或子序列大小； ``num_heads`` 表示transformer模块中的头数； ``head_size`` 表示每个头的维度。
- ``dst``：Vector<Tensor>类型，必选。输入维度 ``[num_layers<num_blocks, num_heads, block_size, head_size>]``。其中每个参数含义与 ``src`` 一致。
- ``block_mapping``：Map<int64, vector<int64>>类型，必选。第一个int64表示需要copy的源BLOCK位置（小于num_blocks），第二个int64表示将源BLOCK copy到的目标位置（小于num_blocks）， ``vector`` 表示一个源BLOCK可以copy到多个目标位置。

**返回值说明**

- dst，原位操作。

**规格限制**

- 数据类型支持：

  - src：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - dst：FP16、FP32、BF16、INT8、UINT8、INT16、INT32、INT64
  - block_mapping：INT64

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_swap_blocks.py`` 。

.. _torch_mlu_ops.weight_only_quant_matmul:

torch_mlu_ops.weight_only_quant_matmul
------------------------------------------

.. code::

  torch_mlu_ops.weight_only_quant_matmul(a,b,scale,zero,bias,c,act_mode,quant_bit_size,alpha,beta,use_hp_active)-> torch.Tensor

**功能描述**

实现仅权重量化功能的GEMM算子。

**公式说明**

* weight_only_quant_matmul标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     ab = matmul(a, b) * scale + bias\\
     acted = active(ab)\\
     output = alpha * acted + beta * c
     \end{array}


* 公式中各参数描述如下：

  - ``a`` ：矩阵乘运算的a矩阵。
  - ``b`` ：矩阵乘运算的b矩阵。
  - ``scale`` : 权重的分channel scale系数。
  - ``bias`` ：偏置。
  - ``alpha`` ：浮点标量，默认1.0f。
  - ``beta`` ：浮点标量，默认1.0f。
  - ``c`` ：第三个elemenwise加c矩阵。
  - ``ouput`` :GEMM返回值。

* 公式中各函数描述如下：
  - ``matmul`` :矩阵乘操作。
  - ``active`` :激活操作。

**参数说明**

- ``a``：Tensor类型，必选。GEMM算子的输入。维度 ``[M, K]``。
- ``b``：Tensor类型，必选。GEMM算子的权重。当 ``quant_bit_size = 8`` ，维度 ``[N, K]``，当 ``quant_bit_size = 8`` ，维度 ``[N, K//2]``。
- ``scale``：Tensor类型，可选。当group_wise量化时，维度为 ``[N, G]``，G表示Group_num，Group_size为K//Group_num，Group_size必须为64、128、256或512，类型必须与a相同。当per_channel量化时，维度为 ``[N]``，类型必须是float，
- ``zero``：Tensor类型，目前不开放，设置为None。
- ``bias``：Tensor类型，可选。GEMM算子的偏置。维度 ``[N]``。
- ``c`` ： Tensor类型，可选。维度 ``[M, N]``。
- ``act_mode`` : string类型，可选的激活函数："relu"、"gelu"、"silu"、"none"。
- ``quant_bit_size`` : int类型，权重 ``b`` 的位宽。
- ``alpha`` ： float类型。默认值为1.0。
- ``beta`` ： float类型。默认值为1.0。
- ``use_hp_active`` ： bool类型，激活的计算方式，默认值为False。


**返回值说明**

- ``ouput``: Tensor类型, 维度 ``[M, N]``。

**规格限制**


- 数据类型支持：

  - a：HALF、FP32、BF16、INT8
  - b: INT8
  - scale：当使用group_wise量化时，类型与a相同，使用per_channel量化时，类型为FLOAT32
  - bias：HALF、FP32、BF16
  - c：HALF、FP32、BF16
  - output：HALF、FP32、BF16

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_quant_ffn.py`` 。

.. _torch_mlu_ops.moe_active:

torch_mlu_ops.moe_active
------------------------------------------

.. code::

  torch_mlu_ops.moe_active(input, act_mode, is_gated, output, bias, cusumt_token_count, start_expert_id, expert_size)-> torch.Tensor

**功能描述**

激活算子。

**公式说明**

* 公式如下：

无。

**参数说明**

- ``input``：Tensor类型，必选。输入维度 ``[..., C]``。 其中 ``C`` 为通道的个数，在最低维。
- ``act_mode``：string类型。可以是 ``silu`` 或者 ``gelu`` 。
- ``is_gated``：Bool类型。当 ``is_gated=true`` 时，为gated激活。
- ``output``：Tensor类型，可选。当 ``is_gated=true`` 时，维度为 ``[..., C//2]``，否则与input形状一致。
- ``bias``：Tensor类型，可选。输入维度是 ``[expert_num, C]`` 。其中 ``expert_num`` 为专家的总数。
- ``cusum_token_count``：Tensor类型，可选，默认值为None。token_count的前缀和，维度 ``[num_expert + 1]`` 。
- ``start_expert_id``：Int32类型，必选。表示EP模型下起始专家的位置信息。
- ``expert_size``：Int32类型，必选。表示EP模型下实际使用的专家数量。

**返回值说明**

-  ``output``：激活后返回的tensor。

**规格限制**

- 数据类型支持：

  - input：FP32、FP16、BF16
  - output：FP32、FP16、BF16
  - bias：FP32、FP16、BF16
  - cusumt_token_count：INT32
  - start_expert_id：INT32
  - expert_size：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_add_bias_activation.py`` 。

.. _torch_mlu_ops.moe_cast_gating:

torch_mlu_ops.moe_cast_gating
------------------------------------------

.. code::

  torch_mlu_ops.moe_cast_gating(input, weight) -> torch.Tensor

**功能描述**

将input的数据类型强转成float类型，再和weight进行矩阵乘操作。

**公式说明**

* 公式如下：

无。

**参数说明**

- ``input``：Tensor类型，必选。输入维度 ``[..., hidden_size]``。 其中 ``hidden_size`` 表示各个token的特征信息。 该输入必须连续。
- ``weight``：Tensor类型，必选。输入维度 ``[expert_num, hidden_size]`` 。 其中 ``expert_num`` 表示专家数量。 该输入必须连续。

**返回值说明**

-  ``output``：Tensor类型。维度为 ``[..., expert_num]``。

**规格限制**

- 数据类型支持：

  - input：FP16、BF16
  - weight：FP32
  - output：FP32

- 规模支持：

  - hidden_size：小于或等于16384
  - expert_num：小于或等于128

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_moe_cast_gating.py`` 。

.. _torch_mlu_ops.update_out_and_lse:

torch_mlu_ops.update_out_and_lse
------------------------------------------

.. code::

  torch_mlu_ops.update_out_and_lse(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)-> None

**功能描述**

根据seq_offset, cu_seqs, block_cu_seqs，将block_out和block_lse的结果更新到out和lse中。

**公式说明**

* 公式如下：
  .. math::

     \begin{array}{lcl}
     out = out - sigmoid(block_lse - lse) * (out - block_out)\\
     lse = lse - logsigmoid(lse - block_lse)\\
     \end{array}

**参数说明**

- ``out``：Tensor类型，必选。输入维度 pad模式： ``[batch, max_seq_len, head_num, head_size]`` pack模式： ``[total_seq_lens, head_num, head_size]``。
- ``lse``：Tensor类型，必选。输入维度 pad模式： ``[batch, head_num, max_seq_len]`` 不支持pack模式。
- ``block_out``：Tensor类型，必选。输入维度 pad模式： ``[batch, block_seq_len, head_num, head_size]`` pack模式： ``[total_block_seq_lens, head_num, head_size]``。
- ``block_lse``：Tensor类型，必选。输入维度 pad模式： ``[batch, head_num, block_seq_len]`` 不支持pack模式。
- ``seq_offset``：Tensor类型，可选。输入维度是 ``[batch]`` 。
- ``cu_seqs``：Tensor类型，可选。输入维度是 ``[batch + 1]`` ，仅限pack模式。
- ``block_cu_seqs``：Tensor类型，可选。输入维度是 ``[batch + 1]`` ，仅限pack模式。

**返回值说明**

无返回值。

**规格限制**

- 数据类型支持：

  - out：FP32、FP16、BF16
  - lse：FP32、FP16、BF16
  - block_out：FP32、FP16、BF16
  - block_lse：FP32、FP16、BF16
  - seq_offset：INT32
  - cu_seqs：INT32
  - block_cu_seqs：INT32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_update_out_and_lse.py`` 。

.. _torch_mlu_ops.single_query_mixed_cached_kv_attn:

torch_mlu_ops.single_query_mixed_cached_kv_attn
-----------------------------------------------------------

.. code::

  torch_mlu_ops.single_query_mixed_cached_kv_attn(q_ori,k_cache_lp,v_cache_lp,k_cache_hp,v_cache_hp,out,block_tables_lp,block_tables_hp,context_lens_lp,context_lens_hp,k_cache_quant_scale_lp,v_cache_quant_scale_lp,k_cache_quant_scale_hp,v_cache_quant_scale_hp,alibi_slopes,max_context_len_lp,max_context_len_hp,softmax_scale,return_lse,kv_cache_quant_bit_size,kv_cache_quant_bit_size_lp,kv_cache_quant_bit_size_hp)->Tensor

**功能描述**

Decoder阶段的context_seq分段量化，采用低精度与高精度融合的Single Query Cached KV Attention算子。

**公式说明**

* context_seq分段量化的融合的Single Query Cached KV Attention标准结构公式如下：

  .. math::

     \begin{array}{lcl}
     lse1, attn\_out1 = single\_query\_cached\_kv\_attn(q\_ori, k\_cache\_lp, v\_cache\_lp, block\_tables\_lp,...)\\
     lse2, attn\_out2 = single\_query\_cached\_kv\_attn(q\_ori, k\_cache\_hp, v\_cache\_hp, block\_tables\_hp,...)\\
     output = update\_out\_and\_lse(lse1, attn\_out1, lse2, attn\_out2)\\
     \end{array}

* 公式中各参数描述如下：

  - ``q_ori`` ：Single Query Cached KV Attention模块的输入。
  - ``attn_out1``、 ``attn_out1``  ：两段Single Query Cached KV Attention的输出。
  - ``lse1``、 ``lse2`` ：两段Single Query Cached KV Attention的lse输出。
  - ``output``：最终的拼接输出。

* 公式中各函数描述如下：

  - ``single_query_cached_kv_attn`` ：Single Query Cached KV Attention算子。
  - ``update_out_and_lse``： update_out_and_lse算子。

**参数说明**

- ``q_ori``：Tensor类型，必选。算子的输入。维度 ``[batch, seq_q, q_head_num, head_size]``。
- ``k_cache_lp``：Tensor类型，必选。算子的输入, 第一段低精度context对应的cache。维度 ``[num_blocks_lp, kv_head_num, block_size_lp, head_size]``。
- ``v_cache_lp``：Tensor类型，必选。算子的输入, 第一段低精度context对应的cache。维度 ``[num_blocks_lp, kv_head_num, block_size_lp, head_size]``。
- ``k_cache_hp``：Tensor类型，必选。算子的输入, 第二段高精度context对应的cache。维度 ``[num_blocks_hp, kv_head_num, block_size_hp, head_size]``。
- ``v_cache_hp``：Tensor类型，必选。算子的输入, 第二段高精度context对应的cache。维度 ``[num_blocks_hp, kv_head_num, block_size_hp, head_size]``。
- ``out``：Tensor类型，可选。维度同input。
- ``block_tables_lp``：Tensor类型，必选。低精度部分block的id与kv_cache之间的映射关系。维度 ``[bs, max_num_blocks_per_seq_lp]``，非paged attention时，max_num_blocks_per_seq_lp为1。
- ``block_tables_hp``：Tensor类型，必选。高精度部分block的id与kv_cache之间的映射关系。维度 ``[bs, max_num_blocks_per_seq_hp]``，非paged attention时，max_num_blocks_per_seq_hp为1。
- ``context_lens_lp``：Tensor类型，必选。记录低精度部分每个batch的实际seq长度，维度 ``[batch]``。
- ``context_lens_hp``：Tensor类型，必选。记录高精度部分每个batch的实际seq长度，维度 ``[batch]``。
- ``k_cache_quant_scale_lp``：Tensor类型，可选。k_cache_lp&v_cache_lp为量化类型时k_cache_lp的量化参数，维度 ``[num_blocks_lp, kv_head_num, block_size_lp]``表示per_token量化、 ``[num_blocks_lp, kv_head_num, block_size_lp, 1]`` 表示带group的per_token量化，但当前group_num必须为1， 以及 ``[kv_head_num, head_size]`` 表示per_channel量化。
- ``v_cache_quant_scale_lp``：Tensor类型，可选。k_cache_lp&v_cache_lp为量化类型时v_cache_lp的量化参数，维度同 ``k_cache_quant_scale``。
- ``k_cache_quant_scale_hp``：Tensor类型，可选。k_cache_hp&v_cache_hp为量化类型时k_cache_hp的量化参数，维度 ``[num_blocks_hp, kv_head_num, block_size_hp]``表示per_token量化、 ``[num_blocks_hp, kv_head_num, block_size_hp, 1]`` 表示带group的per_token量化，但当前group_num必须为1， 以及 ``[kv_head_num, head_size]`` 表示per_channel量化。
- ``v_cache_quant_scale_hp``：Tensor类型，可选。k_cache_hp&v_cache_hp为量化类型时v_cache_hp的量化参数，维度同 ``k_cache_quant_scale``。
- ``alibi_slope``：Tensor类型，可选。是一种相对位置编码方法，维度 ``[batch, q_head_num]``，类型必须是float32。
- ``max_context_len_lp``：Int64类型，必选。算子低精度部分的最大seq长度。
- ``max_context_len_hp``：Int64类型，必选。算子高精度部分的最大seq长度。
- ``softmax_scale``：Float32类型，必选。softmax之前的系数。
- ``return_lse``：Bool类型，标记是否返回lse tensor。
- ``kv_cache_quant_bit_size_lp``：INT32类型，低精度部分KV cache量化的位宽，可选8（表示int8）、4（表示int4）和-1（表示和input一样的数据类型）。
- ``kv_cache_quant_bit_size_hp``：INT32类型，高精度部分KV cache量化的位宽，可选8（表示int8）、4（表示int4）和-1（表示和input一样的数据类型）。

**返回值说明**

- ``output``：Tensor类型，可选。算子的输出，当out不为空时有此输出，维度 ``[batch, seq_q, q_head_num, head_size]``。
- ``output_lse``：Tensor类型，可选。算子的lse输出，当 ``return_lse`` 为True时有此输出，维度 ``[batch, q_head_num, seq_q]``。

**规格限制**

- 数据类型支持：

  - q_ori：FP16、BF16、FP32
  - k_cache_lp：FP16、BF16、FP32、INT8、int4
  - v_cache_lp：FP16、BF16、FP32、INT8、int4
  - k_cache_hp：FP16、BF16、FP32、INT8、int4
  - v_cache_hp：FP16、BF16、FP32、INT8、int4
  - block_tables_lp：INT32
  - block_tables_hp：INT32
  - context_lens_lp：INT32
  - context_lens_hp：INT32
  - k_cache_quant_scale_lp：FP32
  - v_cache_quant_scale_lp：FP32
  - k_cache_quant_scale_hp：FP32
  - v_cache_quant_scale_hp：FP32
  - alibi_slope：FP32
  - max_context_len_lp：INT64
  - max_context_len_hp：INT64
  - softmax_scale：FP32
  - output：FP16、BF16、FP32
  - output_lse：FP32

**使用示例**

请参考 ``torch_mlu_ops/tests/ops_pytest/test_single_query_mixed_cached_kv_attn.py`` 。
