API功能： 根据输入词频[object Object]、[object Object]/[object Object]采样参数、随机采样权重分布[object Object]，进行topK-topP-Sample采样计算，输出每个batch的最大词频[object Object]，以及topK-topP采样后的词频分布[object Object]。

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

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

topK采样

1. 按分段长度v采用分段topK归并排序，用{s-1}块的topK对当前{s}块的输入进行预筛选，渐进更新单batch的topK，减少冗余数据和计算。
2. top_k[batch]对应当前batch采样的k值，有效范围为1≤top_k[batch]≤min(voc_size[batch], 1024)，如果top_k[batch]超出有效范围，则视为跳过当前batch的topK采样阶段，也同样会则跳过当前batch的排序，将输入logits[batch]直接传入下一模块。[object Object]

• 具体计算流程如下所示：
• 对当前batch分割为若干子段，滚动计算top_k_value[b]：$$top\_k\_value[b] = {Max(top\_k[b])}_{s=1}^{\left \lceil \frac{S}{v} \right \rceil }\left \{ top\_k\_value[b]\left \{s-1 \right \} \cup \left \{ logits[b][v] \ge top\_k\_min[b][s-1] \right \} \right \}\\ Card(top\_k\_value[b])=top\_k[b]$$ 其中：$$top\_k\_min[b][s] = Min(top\_k\_value[b]\left \{ s \right \})$$ v表示预设的滚动topK时的固定分段长度：$$v=8*1024$$
• 生成需要过滤的mask：$$sorted\_value[b] = sort(top\_k\_value[b], descendant)$$ $$top\_k\_mask[b] = sorted\_value[b]<Min(top\_k\_value[b])$$
• 将小于阈值的部分通过mask置为-inf：$$sorted\_value[b][v]= \begin{cases} -inf & \text{top\_k\_mask[b][v]=true} \\ sorted\_value[b][v] & \text{top\_k\_mask[b][v]=false} & \end{cases}$$
• 通过softmax将经过topK过滤后的[object Object]按最后一轴转换为概率分布：$$probs\_value[b]=sorted\_value[b].softmax (dim=-1)$$
• 按最后一轴计算累积概率（从最小的概率开始累加）：$$probs\_sum[b]=probs\_value[b].cumsum (dim=-1)$$ topP采样
• 如果前序topK采样已有排序输出结果，则根据topK采样输出计算累积词频，并根据topP截断采样：$$top\_p\_mask[b] = probs\_sum[b][*] < top\_p[b]$$
• 如果topK采样被跳过，则先对输入logits[b]进行softmax处理：$$logits\_value[b] = logits[b].softmax(dim=-1)$$
• 尝试使用[object Object]，对[object Object]进行滚动排序，获取计算topP采样的mask：$$top\_p\_value[b] = {Max(top\_k\_guess)}_{s=1}^{\left \lceil \frac{S}{v} \right \rceil }\left \{ top\_p\_value[b]\left \{s-1 \right \} \cup \left \{ logits\_value[b][v] \ge top\_k\_min[b][s-1] \right \} \right \}$$
• 如果在访问到logits_value[b]的第1e4个元素之前，如下条件得到满足，则视为[object Object]成功：$$\sum^{top\_k\_guess}(top\_p\_value[b]) \ge top\_p[b]\\ top\_p\_mask[b][Index(top\_p\_value[b])] = false$$
• 如果[object Object]失败，则对当前序logits_value[b]进行全排序和cumsum，按top_p[b]截断采样：$$sorted\_logits[b] = sort(logits\_value[b], descendant) \\ probs\_sum[b]=sorted\_logits[b].cumsum (dim=-1) \\ top\_p\_mask[b] = (probs\_sum[b] - sorted\_logits[b])>top\_p[b]$$
• 将需要过滤的位置设置为-inf，得到sorted_value[b][v]：$$sorted\_value[b][v] = \begin{cases} -inf& \text{top\_p\_mask[b][v]=true}\\sorted\_value[b][v]& \text{top\_p\_mask[b][v]=false}\end{cases}$$
• 取过滤后sorted_value[b][v]每行中前topK个元素，查找这些元素在输入中的原始索引，整合为logits_idx：$$logits\_idx[b][v] = Index(sorted\_value[b][v] \in logits)$$ 指数采样（Sample）
• 如果[object Object]设置为True，则根据logits_idx，选取采样后结果输出到[object Object]：$$logits\_top\_kp\_select[b][logits\_idx[b][v]]=sorted\_value[b][v]$$
• 对sorted_value进行指数分布采样：$$probs = softmax(sorted\_value)$$ $$probs\_opt = \frac{probs}{q + eps}$$
• 从probs_opt中取出每个batch的最大元素，从logits_idx中gather相应元素的输入索引，作为输出[object Object]：$$logits\_select\_idx[b] = logits\_idx[b][argmax(probs\_opt[b][:])]$$ 其中0≤b<$sorted\_value.size(0)$，0≤v<$sorted\_value.size(1)$。