aclnnTopKTopPSample

产品支持情况

产品	是否支持
[object Object]Atlas A3 训练系列产品/Atlas A3 推理系列产品[object Object]	√
[object Object]Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件[object Object]	√
[object Object]Atlas 200I/500 A2 推理产品[object Object]	×
[object Object]Atlas 推理系列产品 [object Object]	×
[object Object]Atlas 训练系列产品[object Object]	×

功能说明

• 算子功能： 根据输入词频logits、topK/topP采样参数、随机采样权重分布q，进行topK-topP-sample采样计算，输出每个batch的最大词频logitsSelectIdx，以及topK-topP采样后的词频分布logitsTopKPSelect。

  算子包含三个可单独使能，但上下游处理关系保持不变的采样算法（从原始输入到最终输出）：TopK采样、TopP采样、指数采样（本文档中Sample所指）。它们可以构成八种计算场景。如下表所示：

  计算场景	TopK采样	TopP采样	指数分布采样	备注
  Softmax-Argmax采样	×	×	×	对输入logits按每个batch，取SoftMax后取最大结果
  topK采样	√	×	×
  topP采样	×	√	×
  Sample采样	×	×	√
  topK-topP采样	√	√	×
  topK-Sample采样	√	×	√
  topP-Sample采样	×	√	√
  topK-topP-Sample采样	√	√	√

• 计算公式： 输入logits为大小为[batch, voc_size]的词频表，其中每个batch对应一条输入序列，而voc_size则是约定每个batch的统一长度。[object Object] logits中的每一行logits[batch][:]根据相应的topK[batch]、topP[batch]、q[batch, :]，执行不同的计算场景。[object Object] 下述公式中使用b和v来分别表示batch和voc_size方向上的索引。

  TopK采样

  1. 按分段长度v采用分段topk归并排序，用{s-1}块的topK对当前{s}块的输入进行预筛选，渐进更新单batch的topK，减少冗余数据和计算。
  2. topK[batch]对应当前batch采样的k值，有效范围为1≤topK[batch]≤min(voc_size[batch], 1024)，如果top[k]超出有效范围，则视为跳过当前batch的topK采样阶段，也同样会则跳过当前batch的排序，将输入logits[batch]直接传入下一模块。
  • 对当前batch分割为若干子段，滚动计算topKValue[b]：
  $$topKValue[b] = {Max(topK[b])}_{s=1}^{\left \lceil \frac{S}{v} \right \rceil }\left \{ topKValue[b]\left \{s-1 \right \} \cup \left \{ logits[b][v] \ge topKMin[b][s-1] \right \} \right \}\\ Card(topKValue[b])=topK[b]$$

  其中：

  $$topKMin[b][s] = Min(topKValue[b]\left \{ s \right \})$$

  v表示预设的滚动topK时固定的分段长度：

  $$v=8*1024$$
  • 生成需要过滤的mask
  $$sortedValue[b] = sort(topKValue[b], descendant)$$ $$topKMask[b] = sortedValue[b]<Min(topKValue[b])$$
  • 将小于阈值的部分通过mask置为-inf
  $$sortedValue[b][v]= \begin{cases} -inf & \text{topKMask[b][v]=true} \\ sortedValue[b][v] & \text{topKMask[b][v]=false} & \end{cases}$$
  • 通过softmax将经过topK过滤后的logits按最后一轴转换为概率分布
  $$probsValue[b]=sortedValue[b].softmax (dim=-1)$$
  • 按最后一轴计算累积概率（从最小的概率开始累加）
  $$probsSum[b]=probsValue[b].cumsum (dim=-1)$$

  TopP采样

  • 如果前序topK采样已有排序输出结果，则根据topK采样输出计算累积词频，并根据topP截断采样：

    $$topPMask[b] = probsSum[b][*] < topP[b]$$

  • 如果topK采样被跳过，则先对输入logits[b]进行softmax处理：

    $$logitsValue[b] = logits[b].softmax(dim=-1)$$

  • 尝试使用topKGuess，对logits进行滚动排序，获取计算topP的mask：

    $$topPValue[b] = {Max(topKGuess)}_{s=1}^{\left \lceil \frac{S}{v} \right \rceil }\left \{ topPValue[b]\left \{s-1 \right \} \cup \left \{ logitsValue[b][v] \ge topKMin[b][s-1] \right \} \right \}$$

  • 如果在访问到logitsValue[b]的第1e4个元素之前，下条件得到满足，则视为topKGuess成功，

    $$\sum^{topKGuess}(topPValue[b]) \ge topP[b]\\ topPMask[b][Index(topPValue[b])] = false$$

  • 如果topKGuess失败，则对当前序logitsValue[b]进行全排序和cumsum，按topP[b]截断采样：

    $$sortedLogits[b] = sort(logitsValue[b], descendant) \\ probsSum[b]=sortedLogits[b].cumsum (dim=-1) \\ topPMask[b] = (probsSum[b] - sortedLogits[b])>topP[b]$$

  • 将需要过滤的位置设置为-inf，得到sortedValue[b][v]：

    $$sortedValue[b][v] = \begin{cases} -inf& \text{topPMask[b][v]=true}\\sortedValue[b][v]& \text{topPMask[b][v]=false}\end{cases}$$

    取过滤后sortedValue[b][v]每行中前topK个元素，查找这些元素在输入中的原始索引，整合为logitsIdx:

    $$logitsIdx[b][v] = Index(sortedValue[b][v] \in logits)$$

  指数采样（Sample）

  • 如果isNeedLogits=true，则根据logitsIdx，选取采样后结果输出到logitsTopKPSelect：
  $$logitsTopKPSelect[b][logitsIdx[b][v]]=sortedValue[b][v]$$
  • 对logitsSort进行指数分布采样：$$probs = softmax(logitsSort)$$ $$probsOpt = \frac{probs}{q + eps}$$
  • 从probsOpt中取出每个batch的最大元素，从logitsIdx中gather相应元素的输入索引，作为输出logitsSelectIdx：$$logitsSelectIdx[b] = logitsIdx[b][argmax(probsOpt[b][:])]$$

  其中0≤b<sortedValue.size(0)，0≤v<sortedValue.size(1)。

函数原型

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

• aclnnStatus aclnnTopKTopPSampleGetWorkspaceSize(const aclTensor *logits, const aclTensor *topK, const aclTensor *topP, const aclTensor *q, const aclTensor *logitsSelectIdx, const aclTensor *logitsTopKPSelect, float eps, uint32_t isNeedLogits, uint32_t topKGuess, uint64_t *workspaceSize, aclOpExecutor **executor)
• aclnnStatus aclnnTopKTopPSample(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)

aclnnTopKtopPSampleGetWorkspaceSize

• 参数说明：

  • logits(aclTensor*，计算输入)：Device侧的aclTensor，表示待采样的输入词频，词频索引固定为最后一维。目前支持2维；数据类型支持FLOAT16和BFLOAT16。支持ND，支持非连续Tensor，必须为非空指针。
  • topK(aclTensor*，计算输入)：Device侧的aclTensor，表示每个batch采样的k值。目前支持1维，维度需与logits前n-1维一致；数据类型支持INT32。支持ND，支持非连续Tensor，必须为非空指针。
  • topP(aclTensor*，计算输入)：Device侧的aclTensor，表示每个batch采样的p值。目前支持1维，维度需与logits前n-1维一致；数据类型支持FLOAT16和BFLOAT16，数据类型必须与logits保持一致。支持ND，支持非连续Tensor，必须为非空指针。
  • q(aclTensor*，计算输入)：Device侧的aclTensor，表示topK-topP采样输出的指数采样矩阵。目前支持2维，维度与尺寸需要与logits保持一致；数据类型支持FLOAT32。支持ND，支持非连续Tensor，必须为非空指针。
  • eps（float，入参）：在softmax和权重采样中防止除零，建议设置为1e-8。
  • isNeedLogits（uint32_t，入参）：控制logitsTopKPselect的输出条件，建议设置为0。
  • topKGuess（uint32_t，入参）：表示每个batch在尝试topP部分遍历采样logits时的候选logits大小，必须为正整数。
  • logitsSelectIdx(aclTensor*，计算输出)：Device侧的aclTensor，表示经过topK-topP-sample计算流程后，每个batch中词频最大元素max(probsOpt[batch, :])在输入logits中的位置索引。数据类型支持INT64，数据格式支持ND。必须为非空指针。
  • logitsTopKPSelect(aclTensor*，计算输出)：Device侧的aclTensor，表示经过topK-topP计算流程后，输入logits中剩余未被过滤的logits。数据类型支持FLOAT32，数据格式支持ND。isNeedLogits=0时不使能，输出空指针；当isNeedLogits=1时使能，输出尺寸必然为[batch, voc_size]。
  • workspaceSize(uint64_t*, 出参)：返回需要在Device侧申请的workspace大小。
  • executor(aclOpExecutor**, 出参)：返回op执行器，包含了算子计算流程。

• 返回值：

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

  [object Object]

aclnnTopKtopPSample

• 参数说明：

  • workspace(void*, 入参)：在Device侧申请的workspace内存地址。
  • workspaceSize(uint64_t, 入参)：在Device侧申请的workspace大小，由第一段接口aclnnTopKtopPSampleGetWorkspaceSize获取。
  • executor(aclOpExecutor*, 入参)：op执行器，包含了算子计算流程。
  • stream(aclrtStream, 入参)：指定执行任务的AscendCL Stream流。

• 返回值：

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

约束说明

• 对于所有参数，它们的尺寸必须满足，batch>0，0<vocSize<=2^20。
• logits、q、logitsTopKPselect的尺寸和维度必须完全一致。
• logits、topK、topP、logitsSelectIdx除最后一维以外的所有维度必须顺序和大小完全一致。目前logits只能是2维，topK、topP、logitsSelectIdx必须是1维非空Tensor。logits、topK、topP不允许空Tensor作为输入，如需跳过相应模块，需按相应规则设置输入。
• 如果需要单独跳过topK模块，请传入[batch, 1]大小的Tensor，并使每个元素均为无效值。
• 如果1024<topK[batch]<vocSize[batch]，则视为选择当前batch的全部有效元素并跳过topK采样。
• 如果需要单独跳过topP模块，请传入[batch, 1]大小的Tensor，并使每个元素均≥1。
• 如果需要单独跳过sample模块，传入q=nullptr即可；如需使用sample模块，则必须传入尺寸为[batch, vocSize]的Tensor。

调用示例

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