aclnnApplyTopKTopP

产品支持情况

[object Object]undefined

功能说明

• 接口功能：对原始输入logits进行top-k和top-p采样过滤。

• 计算公式：

  • 对输入logits按最后一轴进行升序排序，得到对应的排序结果sortedValue和sortedIndices。

  $$sortedValue, sortedIndices = sort(logits, dim=-1, descending=false, stable=true)$$

  • 计算保留的阈值（第k大的值）。

  $$topKValue[b][v] = sortedValue[b][sortedValue.size(1) - k[b]]$$

  • 生成top-k需要过滤的mask。

  $$topKMask = sortedValue < topKValue$$

  • 通过topKMask将小于阈值的部分置为-Inf。
  $$sortedValue[b][v] = \begin{cases} -Inf & \text{topKMask[b][v]=true}\\ sortedValue[b][v] & \text{topKMask[b][v]=false} \end{cases}$$
  • 通过softmax将经过top-k过滤后的数据按最后一轴转换为概率分布。

  $$probsValue = softmax(sortedValue, dim=-1)$$

  • 按最后一轴计算累计概率（从最小的概率开始累加）

  $$probsSum = cumsum(probsValue, dim=-1)$$

  • 生成top-p的mask，累计概率小于等于1-p的位置需要过滤掉，并保证每个batch至少保留一个元素。

  $$topPMask[b][v] = probsSum[b][v] <= 1-p[b]$$

  $$topPMask[b][-1] = false$$

  • 通过topPMask将小于阈值的部分置为-Inf。
  $$sortedValue[b][v] = \begin{cases} -Inf & \text{topPMask[b][v]=true}\\ sortedValue[b][v] & \text{topPMask[b][v]=false} \end{cases}$$
  • 将过滤后的结果按sortedIndices还原到原始顺序。

  $$out[b][v] = sortedValue[b][sortedIndices[b][v]]$$

  其中$0 \le b \lt logits.size(0), 0 \le v \lt logits.size(1)$。

函数原型

每个算子分为，必须先调用“aclnnApplyTopKTopPGetWorkspaceSize”接口获取计算所需workspace大小以及包含了算子计算流程的执行器，再调用“aclnnApplyTopKTopP”接口执行计算。

aclnnApplyTopKTopPGetWorkspaceSize

• 参数说明：

  [object Object]

• 返回值：

  aclnnStatus：返回状态码，具体参见。

  第一段接口会完成入参校验，出现以下场景时报错：

  [object Object]

aclnnApplyTopKTopP

• 参数说明：

  [object Object]

• 返回值：

  aclnnStatus：返回状态码，具体参见。

约束说明

• 确定性计算：
  • aclnnApplyTopKTopP默认确定性实现。

调用示例

示例代码如下，仅供参考，具体编译和执行过程请参考。