[object Object][object Object][object Object]undefined[object Object]

• Description: Performs topK-topP sampling computation based on the input logits, topK/topP sampling parameters, and random sampling weight distribution q, and outputs the maximum word frequency logitsSelectIdx of each batch and the word frequency distribution logitsTopKPSelect after topK-topP sampling.

  The operator contains three sampling algorithms that can be enabled separately with the upstream and downstream processing relationships remain unchanged (from the original input to the final output): TopK sampling, TopP sampling, and exponential sampling (referred to as Sample in this document). They can form eight computing scenarios. 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] Based on topK[batch], topP[batch], and q[batch,:], different computations are executed for each logits[batch][:] row in logits.[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

  1. 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.
  2. topK[batch] indicates the k value of the current batch. The valid range is 1 ≤ topK[batch] ≤ min(voc_size[batch], 1024). If top[k] is out of the valid range, the topK sampling phase of the current batch is skipped, so does the sorting of the current batch. The input logits[batch] is directly passed to the next module.
  • Divides the current batch into several sub-segments and calculates topKValue[b] in rolling mode.
  $$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]$$

  Where:

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

  v indicates the preset fixed segment length during the rolling topK operation.

  $$v=8*1024$$
  • Generates the mask to be filtered.
  $$sortedValue[b] = sort(topKValue[b], descendant)$$ $$topKMask[b] = sortedValue[b]<Min(topKValue[b])$$
  • Sets the part that is less than the threshold to -Inf using mask.
  $$sortedValue[b][v]= \begin{cases} -Inf & \text{topKMask[b][v]=true} \\ sortedValue[b][v] & \text{topKMask[b][v]=false} & \end{cases}$$
  • Converts the logits filtered by topK into probability distribution along the last axis using softmax.
  $$probsValue[b]=sortedValue[b].softmax (dim=-1)$$
  • Computes the cumulative probability along the last axis (starting cumulation from the smallest probability).
  $$probsSum[b]=probsValue[b].cumsum (dim=-1)$$

  TopP sampling

  • If the previous topK sampling has a sorted output result, computes the cumulative word frequency based on the topK sampling output, and performs truncated sampling based on the topP:

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

  • If topK sampling is skipped, performs softmax on the input logits[b] first:

  $$logitsValue[b] = logits[b].softmax(dim=-1)$$
  • Attempts to use topKGuess to perform rolling sorting on logits to obtain the mask for computing topP:
  $$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 \}$$
  • If the following condition is met before the 1e4th element of logitsValue[b] is accessed, topKGuess is considered successful:
  $$\sum^{topKGuess}(topPValue[b]) \ge topP[b]\\ topPMask[b][Index(topPValue[b])] = false$$
  • If topKGuess fails, performs full sorting and cumsum on the current logitsValue[b], and performs truncated sampling based on topP[b]:
  $$sortedLogits[b] = sort(logitsValue[b], descendant) \\ probsSum[b]=sortedLogits[b].cumsum (dim=-1) \\ topPMask[b] = (probsSum[b] - sortedLogits[b])>topP[b]$$
  • Sets the positions to be filtered to -Inf to obtain sortedValue[b][v]:$$sortedValue[b][v] = \begin{cases} -Inf& \text{topPMask[b][v]=true}\\sortedValue[b][v]& \text{topPMask[b][v]=false}\end{cases}$$ Obtains the first topK elements in each row of the filtered sortedValue[b][v], searches for the original indices of these elements in the input, and integrates them into [object Object]:$$logitsIdx[b][v] = Index(sortedValue[b][v] \in logits)$$

  Exponential sampling (Sample)

  • If [object Object], selects the sampling result based on [object Object] and outputs it to [object Object]:
  $$logitsTopKPSelect[b][logitsIdx[b][v]]=sortedValue[b][v]$$

  • Performs exponential distribution sampling on [object Object]:

    $$probs = softmax(logitsSort)$$ $$probsOpt = \frac{probs}{q + eps}$$

  • Obtains the maximum element of each batch from [object Object], performs gather on the input indices of the corresponding elements from [object Object], and uses the gathered result as the output [object Object]:

    $$logitsSelectIdx[b] = logitsIdx[b][argmax(probsOpt[b][:])]$$

  where, 0 ≤ b < sortedValue.size(0) and 0 ≤ v < sortedValue.size(1)

[object Object]

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 process. Then, [object Object] is called to perform computation.

[object Object]

• Parameters:

  [object Object]

• Returns:

  [object Object]: status code. For details, see . The first-phase API implements input parameter verification. The following errors may be thrown:

  [object Object]

[object Object]

• Parameters:

  [object Object]

• Returns:

  [object Object]: status code. For details, see .

[object Object]

• Deterministic compute:
  • aclnnTopKTopPSample defaults to a deterministic implementation.
• For all sampling parameters, their sizes must meet the following requirements: batch > 0 and 0 < vocSize <= 2^20.
• Only non-negative values are valid inputs for topK and topP. Passing 0 or negative values will not cause the sampling to be skipped for the corresponding batch; instead, it causes unexpected errors.
• The sizes of logits, q, and logitsTopKPselect must be the same, so does their dimensions.
• All dimensions except the last dimension of logits, topK, topP, and logitsSelectIdx must be in the same order and have the same size. 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 1024 < topK[batch] < vocSize[batch], all valid elements in the current batch are selected and topK sampling is skipped.
• 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 sample module separately, pass [object Object]. To use the sample module, pass a tensor of size [batch, vocSize].

[object Object]

The following example is for reference only. For details, see .