API function: Performs topK-topP-minP-sample sampling calculation based on the input word frequency logits, topK/topP/minP sampling parameters, and random sampling weight distribution q. If isNeedSampleResult is set to false, the maximum word frequency logitsSelectIdx of each batch and the word frequency distribution logitsTopKPSelect after topK-topP-minP sampling are output. If isNeedSampleResult is set to true, the intermediate calculation results logitsIdx and logitsSortMasked after topK-topP-minP sampling are output. logitsSortMasked is the intermediate result of the word frequency logits after topK-topP-minP sampling calculation, and logitsIdx is the index of logitsSortMasked in logits.
The operator contains four sampling algorithms (from the original input to the final output) that can be enabled separately but the upstream and downstream processing relationships remain unchanged: TopK sampling, TopP sampling, MinP sampling, and exponential sampling (as described in this document). Currently, the following computing scenarios are supported. as follows.
[object Object]undefined
Formula: The input logits are a word frequency table with a size of [batch, voc_size], where each batch corresponds to one input sequence, and voc_size is the uniform length of each batch.[object Object] Each row logits[batch][:] in logits performs different calculation scenarios based on the corresponding topK[batch], topP[batch], minP[batch, :], and q[batch, :].[object Object] In the following formulas, b and v are used to represent the indices in the batch and voc_size directions, respectively.
TopK sampling
- Computes the topK for each segment based on the segment length v, applies merge sort to the topK results, pre-filters the input of the current {s} segments using the topK of the previous {s-1} segments, and gradually updates the topK of a single batch to reduce data redundancy and computation.
- topK[batch] corresponds to the k value sampled in the current batch. The valid range is 1 ≤ topK[batch] ≤ min(voc_size[batch], ks_max). If top[k] is out of the valid range, the topK sampling phase of the current batch is skipped, and the current batch is also skipped. In this case, the input logits[batch] is directly transferred to the next module.
- Divides the current batch into several sub-segments and calculates topKValue[b] in rolling mode.
Where:
v indicates the preset fixed segment length during the rolling topK operation.
The value range of
[object Object]is [1, 1024] with a default value of[object Object], and it must be rounded up to a multiple of 8.- Generates the mask to be filtered.
- Set the part that is less than the threshold to defLogit using the mask.
- defLogit depends on the input_is_logits attribute of the input parameter. This attribute controls the normalization of the input logits and output logits_top_kp_select.
TopP sampling
Based on the inputIsLogits attribute of the input parameter, if the attribute is True, the sorted result is normalized.
The processing strategy of this module varies depending on the value of the input
[object Object]:[object Object]undefined
If regular top-P sampling is performed and the preceding top-K stage has already produced sorted results, the cumulative probabilities are computed based on the top-K output, and sampling is truncated according to
[object Object].If regular top-P sampling is performed but the preceding top-K stage is skipped, the top-P mask is computed.
The positions to be filtered are set to the default invalid value defLogit, and logits_sort is obtained and recorded as sortedValue[b][v].
Obtain the first topK elements in each row of sortedValue[b][v] after filtering, find the original indices of these elements in the input, and integrate them into logits_idx.
Use the truncated sortedValue as logitsSortMasked.
minP sampling
If
[object Object], perform min-P sampling.In other cases:
When
[object Object], only 1 token with the maximum logit value is retained for each batch, and all other positions are filled with[object Object].
Optional output
- If the input parameter IsNeedLogits is set to True, logitsIndexMasked generated after topK-topP-minP joint sampling is used for
[object Object]output. If the top-K, top-P, or min-P sampling stage is skipped, its corresponding mask is set to all ones. - The logitsIndexMasked is used to select the input logits and filter out high-frequency tokens in the input logits as the
[object Object]output.
Subsequent processing
- The input to this stage is
[object Object], which is the combined result of the preceding top-K, top-P, and min-P sampling stages. - Ensure that logitsSortMasked∈(0,1) is used as the input. The inputIsLogits attribute is configured based on the actual input logits. That is: Ensure that There are three modes: None, QSample, and output intermediate results. The modes are controlled by the input parameter constraints isNeedSampleResult and whether q is input.
- None:
- This mode is used when isNeedSampleResult is false and q is not input. In this mode, the maximum element and index are obtained for each batch by using Argmax, and the result is output through gatherOut.
- QSample:
- This mode is used when isNeedSampleResult is false and q is input. In this mode, exponential distribution sampling is performed on probs first.
- Then, generate the output through
[object Object]and[object Object]: - Intermediate result output:
- This mode is used when isNeedSampleResult is true. In this case, the sampled logitsSortMasked and its original index logitsIdx in the input are output.
Each operator has calls. First, [object Object] is called to obtain the workspace size required for computation and the executor that contains the operator computation flow. Then, [object Object] is called to perform computation.
Parameter description:
[object Object]Returns:
[object Object]: status code. For details, see .The first-phase API implements input parameter validation. The following error codes may be returned.
[object Object]
- Deterministic computation:
- The default deterministic implementation of aclnnTopKTopPSampleV2 is used.
- For all sampling parameters, their sizes must meet the following requirements: batch > 0 and 0 < vocSize <= 2^20.
- topK only accepts non-negative values as valid inputs. If 0 or a negative value is passed, the sampling of the corresponding batch is skipped.
- The sizes and dimensions of logits, q, logitsTopKPselect, logitsIdx, and logitsSortMasked must be the same.
- The sizes and dimensions of logits, topK, topP, minPs, logitsSelectIdx, logitsIdx, and logitsSortMasked, except the last dimension, must be the same. Currently, logits can only be two-dimensional, and topK, topP, and logitsSelectIdx must be one-dimensional non-empty tensors. Empty tensors cannot be used as the input of logits, topK, or topP. If the corresponding module needs to be skipped, set the input as required.
- To skip the topK module separately, pass a tensor of size [batch, 1] and set each element to an invalid value.
- If min(ksMaxAligned, 1024)<topK[batch]<vocSize[batch], all valid elements in the current batch are selected and the topK sampling is skipped. ksMaxAligned is the value of ksMax rounded up to the nearest multiple of 8. The value range of ksMax is [1, 1024].
- To skip the topP module separately, pass a tensor of size [batch, 1] and set each element to a value greater than or equal to 1.
- To skip the minP module, pass
[object Object]or a tensor of size [batch, 1] and ensure that each element is less than or equal to 0. - To skip the sample module separately, pass q=nullptr. To use the sample module, pass a tensor of size [batch, vocSize].
- If intermediate results are required, set isNeedSampleResult to true and pass
[object Object]. In this case, logitsSelectIdx is not output.
The following example is for reference only. For details, see .