aclnnTopKTopPSample

Applicable Products

Product	Supported
Ascend 950PR/Ascend 950DT	×
Atlas A3 training products/Atlas A3 inference products	√
Atlas A2 training products/Atlas A2 inference products	√
Atlas 200I/500 A2 inference products	×
Atlas inference products	×
Atlas training products	×

Function

Description:

Performs TopK-TopP-Sample sampling computation based on the input term frequency logits, topK/topP sampling parameters, and the random sampling weight distribution q, and outputs the maximum term frequency logitsSelectIdx for each batch, as well as the term frequency distribution logitsTopKPSelect after TopK-TopP sampling.

The operator includes three sampling algorithms that can be enabled individually but maintain a fixed upstream-downstream processing relationship (from original input to final output): TopK sampling, TopP sampling, and exponential sampling (referred to as Sample in this document). They can form eight compute scenarios, as shown in the following table:

Compute Scenario	TopK Sampling	TopP Sampling	Exponential Sampling	Remarks
Softmax-Argmax sampling	×	×	×	On input logits per each batch, take the maximum result after SoftMax.
TopK sampling	√	×	×	On input logits per each batch, take the top topK[batch] maximum results.
TopP sampling	×	√	×	On input logits per each batch, sort in descending order, and sample the first n results whose cumulative sum is greater than or equal to the topP[batch] value.
Sample sampling	×	×	√	On input logits per each batch, perform SoftMax then divide by q and take the maximum result.
TopK-TopP sampling	√	√	×	On input logits per each batch, perform TopK sampling first, then perform TopP sampling and take the maximum result.
TopK-Sample sampling	√	×	√	On input logits per each batch, perform TopK sampling first, then perform Sample sampling and take the maximum result.
TopP-Sample sampling	×	√	√	On input logits per each batch, perform TopP sampling first, then perform Sample sampling and take the maximum result.
TopK-TopP-Sample sampling	√	√	√	On input logits per each batch, perform TopK sampling first, then perform TopP sampling, and finally perform Sample sampling and take the maximum result.

Formula:

The input logits is a term frequency table of size [batch, voc_size], where each batch corresponds to an input sequence, and voc_size is the uniform length agreed upon for each batch.
Each row logits[batch][:] in logits executes different compute scenarios based on the corresponding topK[batch], topP[batch], and q[batch, :].
In the following formulas, b and v are used to denote indices along the batch and voc_size dimensions, respectively.

TopK Sampling

Perform segmented TopK merge sort by segment length v, using the TopK of block {s-1} to pre-filter the input of the current block {s}, progressively updating the TopK of a single batch to reduce redundant data and computation.
topK[batch] corresponds to the k value sampled for the current batch, with a valid range of 1 ≤ topK[batch] ≤ min(voc_size[batch], 1024). If topK[batch] exceeds the valid range, the TopK sampling phase for the current batch is skipped, and the sorting for the current batch is also skipped, with the input logits[batch] passed directly to the next module.

Split the current batch into several sub-segments and compute topKValue[b] in a rolling manner:

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 fixed segment length used during rolling TopK:

v=8*1024

Generate the mask to be filtered

sortedValue[b] = sort(topKValue[b], descendant)

topKMask[b] = sortedValue[b]<Min(topKValue[b])

Set values below the threshold to -Inf via mask

sortedValue[b][v]= \begin{cases} -Inf & \text{topKMask[b][v]=true} \\ sortedValue[b][v] & \text{topKMask[b][v]=false} & \end{cases}

Convert the TopK-filtered logits into a probability distribution along the last axis via softmax

probsValue[b]=sortedValue[b].softmax (dim=-1)

Computes the cumulative probability along the last axis (accumulating from the smallest probability).

probsSum[b]=probsValue[b].cumsum (dim=-1)

TopP sampling

If the preceding TopK sampling already produces sorted output results, computes the cumulative term frequency based on the TopK sampling output and performs truncation sampling according to TopP:
$topPMask[b] = probsSum[b][*] < topP[b]$
If TopK sampling is skipped, first perform softmax on the input logits[b]:

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

Attempt to use topKGuess to perform rolling sort on logits and 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 accessing the 1e4-th element of logitsValue[b], topKGuess is considered successful:

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

If topKGuess fails, perform full sorting and cumsum on the current logitsValue[b], and truncate sampling by topP[b]:

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

Set the positions to be filtered to -Inf, obtaining sortedValue[b][v]: $sortedValue[b][v] = \begin{cases} -Inf& \text{topPMask[b][v]=true}\\sortedValue[b][v]& \text{topPMask[b][v]=false}\end{cases}$ Take the first topK elements in each row of the filtered sortedValue[b][v], find the original indices of these elements in the input, and consolidate them into logitsIdx: $logitsIdx[b][v] = Index(sortedValue[b][v] \in logits)$

Exponential sampling (Sample)

If isNeedLogits=true, select the sampled result based on logitsIdx and output it to logitsTopKPSelect:

logitsTopKPSelect[b][logitsIdx[b][v]]=sortedValue[b][v]

Perform exponential distribution sampling on logitsSort:
$probs = softmax(logitsSort)$ $probsOpt = \frac{probs}{q + eps}$
Take the maximum element of each batch from probsOpt, and gather the input index of the corresponding element from logitsIdx as the output logitsSelectIdx:
$logitsSelectIdx[b] = logitsIdx[b][argmax(probsOpt[b][:])]$

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

Function Prototype

Each operator is divided into a two-phase API. The aclnnTopKTopPSampleGetWorkspaceSize API must be called first to obtain the workspace size required for computation and the executor that encapsulates the operator computation flow. Then, the aclnnTopKTopPSample API is called to perform the computation.

aclnnStatus aclnnTopKTopPSampleGetWorkspaceSize(
  const aclTensor *logits,
  const aclTensor *topK,
  const aclTensor *topP,
  const aclTensor *q,
  double           eps,
  bool             isNeedLogits,
  int64_t          topKGuess,
  const aclTensor *logitsSelectIdx,
  const aclTensor *logitsTopKPSelect,
  uint64_t        *workspaceSize,
  aclOpExecutor  **executor)

aclnnStatus aclnnTopKTopPSample(
  void           *workspace,
  uint64_t        workspaceSize,
  aclOpExecutor  *executor,
  aclrtStream     stream)

aclnnTopKTopPSampleGetWorkspaceSize

Parameters:

Parameter	Input/Output	Description	Instruction	Data Type	Data Format	Dimension (Shape)	Non-Contiguous Tensor
logits	Input	Indicates the input term frequency to be sampled. The term frequency index is fixed at the last dimension, corresponding to the formula `logits`.	Null tensor not supported.	FLOAT16, BFLOAT16	ND	2	√
topK	Input	Indicates the k value for sampling per each batch. Corresponds to `topK[b]` in the formula.	Null tensor is not supported. The shape must match the first n-1 dimensions of `logits`.	INT32	ND	1	√
topP	Input	Indicates the p value for sampling per batch. Corresponds to `topP[b]` in the corresponding formula.	Null tensor not supported. The shape must match the first n-1 dimensions of `logits`, and the data type must match that of `logits`.	FLOAT16, BFLOAT16	ND	1	√
q	Input	Indicates the exponential sampling matrix output by topK-topP sampling. Corresponds to `q` in the formula.	Null tensor is not supported. The shape must match that of `logits`.	FLOAT32	ND	2	√
eps	Input	Prevents division by zero in softmax and weight sampling. It is recommended to set this to 1e-8.	-	DOUBLE	-	-	-
isNeedLogits	Input	Controls the output condition of logitsTopKPselect. It is recommended to set this to 0.	-	BOOL	-	-	-
topKGuess	Input	Indicates the candidate logits size per batch when attempting to traverse and sample logits in the TopP portion. Must be a positive integer.	-	INT64	-	-	-
logitsSelectIdx	Output	Indicates the position index of the maximum term frequency element max(probsOpt[batch, :]) in the input logits for each batch after the topK-topP-sample computation process.	Null tensor is not supported. The shape must match the first n-1 dimensions of `logits`.	INT64	ND	1	√
logitsTopKPSelect	Output	Indicates the logits that remain unfiltered in the input logits after the TopK-TopP computation process.	Null tensor not supported. The shape must match the first n-1 dimensions of `logits`.	FLOAT32	ND	1	√
workspaceSize	Output	Returns the workspace size to be allocated on the Device side.	-	-	-	-	-
executor	Output	Returns the operator executor, which contains the operator compute flow.	-	-	-	-	-

Return value

aclnnStatus: Return Status. See aclnn Return Code. The first-phase API performs input parameter validation and reports an error in the following scenarios:

Return Code	Error Code	Description
ACLNN_ERR_PARAM_NULLPTR	161001	Any of the input parameters logits, topK, or topP is a null pointer.
ACLNN_ERR_PARAM_INVALID	161002	The data type of logits, topK, topP, or q is not within the supported range.
		The dimension or size of logits and q is inconsistent.
		The dimension of topK or topP is inconsistent with the first n-1 dimensions of logits.
		The data type of logits and topP is inconsistent.

aclnnTopKTopPSample

Parameters:

Parameter	Input/Output	Description
workspace	Input	Device-side-allocated workspace memory address.
workspaceSize	Input	Device-side-allocated workspace size, obtained from the first-phase API aclnnTopKTopPSampleGetWorkspaceSize.
executor	Input	Operator executor that contains the operator compute flow.
stream	Input	Specifies the stream for task execution.

Return value

aclnnStatus: return code. For details, see aclnn Return Code.

Constraints

Deterministic computation:
- aclnnTopKTopPSample defaults to a deterministic implementation.
For all sampling parameters, their sizes must satisfy: batch > 0, 0 < vocSize <= 2^20.
topK and topP accept only non-negative values as valid input. Passing 0 or negative values does not skip the sampling of the corresponding batch; instead, it causes unexpected errors.
The sizes and dimensions of logits, q, and logitsTopKPselect must be exactly the same.
All dimensions of logits, topK, topP, and logitsSelectIdx except the last dimension must be in the same order and of the same size. Currently, logits can only be 2-dimensional, and topK, topP, and logitsSelectIdx must be 1-dimensional non-null tensors. Null tensors are not allowed as input for logits, topK, and topP. To skip the corresponding module, set the input according to the corresponding rules.
To skip the topK module individually, pass in a tensor of size [batch, 1] with each element set to an invalid value.
If 1024 < topK[batch] < vocSize[batch], all valid elements of the current batch are selected and topK sampling is skipped.
To skip the topP module individually, pass in a tensor of size [batch, 1] with each element ≥ 1.
To skip the sample module individually, pass q=nullptr. To use the sample module, a tensor of size [batch, vocSize] must be passed.

Example

The following provides example code for reference only. For details about compilation and running, see Compile and Run Samples.

#include <iostream>
#include <vector>
#include "acl/acl.h"
#include "aclnnop/aclnn_top_k_top_p_sample.h"

#define CHECK_RET(cond, return_expr) \
  do {                               \
    if (!(cond)) {                   \
      return_expr;                   \
    }                                \
  } while (0)

#define LOG_PRINT(message, ...)     \
  do {                              \
    printf(message, ##__VA_ARGS__); \
  } while (0)

int64_t GetShapeSize(const std::vector<int64_t>& shape) {
  int64_t shapeSize = 1;
  for (auto i : shape) {
    shapeSize *= i;
  }
  return shapeSize;
}

int Init(int32_t deviceId, aclrtStream* stream) {
  // Boilerplate for resource initialization.
  auto ret = aclInit(nullptr);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
  ret = aclrtSetDevice(deviceId);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
  ret = aclrtCreateStream(stream);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
  return 0;
}

template <typename T>
int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
                    aclDataType dataType, aclTensor** tensor) {
  auto size = GetShapeSize(shape) * sizeof(T);
  // Call aclrtMalloc to allocate device-side memory.
  auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
  // Call aclrtMemcpy to copy host-side data to the device-side memory.
  ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
  CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);

  // Calculate the strides of a contiguous tensor.
  std::vector<int64_t> strides(shape.size(), 1);
  for (int64_t i = shape.size() - 2; i >= 0; i--) {
    strides[i] = shape[i + 1] * strides[i + 1];
  }

  // Call aclCreateTensor to create an aclTensor.
  *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
                            shape.data(), shape.size(), *deviceAddr);
  return 0;
}

int main() {
    // 1. (Boilerplate) Initialize the device/stream. See the ACL API.
    // Set deviceId based on the actual device.
    int32_t deviceId = 0;
    aclrtStream stream;
    auto ret = Init(deviceId, &stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);

    // 2. Construct the inputs and outputs. Custom construction is required based on the API.
    std::vector<int64_t> logitsShape = {48, 131072};
    std::vector<int64_t> topKPShape = {48};
    long long vocShapeSize = GetShapeSize(logitsShape);
    long long batchShapeSize = GetShapeSize(topKPShape);

    void* logitsDeviceAddr = nullptr;
    void* topKDeviceAddr = nullptr;
    void* topPDeviceAddr = nullptr;
    void* qDeviceAddr = nullptr;
    void* logitsSelectedIdxDeviceAddr = nullptr;
    void* logitsTopKPSelectDeviceAddr = nullptr;

    aclTensor* logits = nullptr;
    aclTensor* topK = nullptr;
    aclTensor* topP = nullptr;
    aclTensor* q = nullptr;
    aclTensor* logitsSelectedIdx = nullptr;
    aclTensor* logitsTopKPSelect = nullptr;
    std::vector<int16_t> logitsHostData(48 * 131072, 1);
    std::vector<int32_t> topKHostData(48, 128);
    std::vector<int16_t> topPHostData(48, 1);
    std::vector<float> qHostData(48 * 131072, 1.0f);

    std::vector<int64_t> logitsSelectedIdxHostData(48, 0);
    std::vector<float> logitsTopKPSelectHostData(48 * 131072, 0);

    float eps=1e-8;
    int64_t isNeedLogits=0;
    int32_t topKGuess=32;
    // Create the logits aclTensor.
    ret = CreateAclTensor(logitsHostData, logitsShape, &logitsDeviceAddr, aclDataType::ACL_BF16, &logits);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create topK aclTensor.
    ret = CreateAclTensor(topKHostData, topKPShape, &topKDeviceAddr, aclDataType::ACL_INT32, &topK);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create topP aclTensor.
    ret = CreateAclTensor(topPHostData, topKPShape, &topPDeviceAddr, aclDataType::ACL_BF16, &topP);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create q aclTensor.
    ret = CreateAclTensor(qHostData, logitsShape, &qDeviceAddr, aclDataType::ACL_FLOAT, &q);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create logitsSelected aclTensor.
    ret = CreateAclTensor(logitsSelectedIdxHostData, topKPShape, &logitsSelectedIdxDeviceAddr, aclDataType::ACL_INT64, &logitsSelectedIdx);
    CHECK_RET(ret == ACL_SUCCESS, return ret);
    // Create logitsTopKPSelect aclTensor.
    ret = CreateAclTensor(logitsTopKPSelectHostData, logitsShape, &logitsTopKPSelectDeviceAddr, aclDataType::ACL_FLOAT, &logitsTopKPSelect);
    CHECK_RET(ret == ACL_SUCCESS, return ret);

    // 3. Call the CANN operator library API. Replace with the specific API name.
    uint64_t workspaceSize = 0;
    aclOpExecutor* executor;
    // Call the first-phase API of aclnnTopKTopPSample.
    ret = aclnnTopKTopPSampleGetWorkspaceSize(logits, topK, topP, q, eps, isNeedLogits, topKGuess, logitsSelectedIdx, logitsTopKPSelect, &workspaceSize, &executor);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTopKTopPSampleGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
    // Apply for device memory based on the workspaceSize calculated by the first-phase API.
    void* workspaceAddr = nullptr;
    if (workspaceSize > 0) {
        ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
        CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
    }
    // Call the second-phase API of aclnnTopKTopPSample.
    ret = aclnnTopKTopPSample(workspaceAddr, workspaceSize, executor, stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnTopKTopPSample failed. ERROR: %d\n", ret); return ret);

    // 4. (Boilerplate) Synchronize and wait for task execution to complete.
    ret = aclrtSynchronizeStream(stream);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);

    // 5. Obtain the output value and copy the result from the device memory to the host memory. Modify based on the specific API definition.
    auto size = GetShapeSize(topKPShape);
    std::vector<float> resultData(size, 0);
    ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), logitsSelectedIdxDeviceAddr,
                        size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
    CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return ret);
    for (int64_t i = 0; i < size; i++) {
        LOG_PRINT("result[%ld] is: %d\n", i, resultData[i]);
    }

    // 6. Release aclTensor. Modify based on the specific API definition.
    aclDestroyTensor(logits);
    aclDestroyTensor(topK);
    aclDestroyTensor(topP);
    aclDestroyTensor(q);
    aclDestroyTensor(logitsSelectedIdx);
    aclDestroyTensor(logitsTopKPSelect);
    // 7. Release device resources. Modify based on the specific API definition.
    aclrtFree(logitsDeviceAddr);
    aclrtFree(topKDeviceAddr);
    aclrtFree(topPDeviceAddr);
    aclrtFree(qDeviceAddr);
    aclrtFree(logitsSelectedIdxDeviceAddr);
    aclrtFree(logitsTopKPSelectDeviceAddr);
    if (workspaceSize > 0) {
        aclrtFree(workspaceAddr);
    }
    aclrtDestroyStream(stream);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
    }