Input and Output

topkToppSamplingType = SINGLE_TOPK_SAMPLING (Non-batch Random Seed and topk Sampling)

Parameter

Dimension

Data Type

Format

Description

probs

[batch, voc_size]

float16/bf16

ND

Input, vocabulary probability.

topp

[batch, 1]

float16/bf16

ND

Input topp. The value of batch must be the same as that of probs.

sampled_indices

[batch, 1]

int32

ND

Output, sampled idx.

sampled_probs

[batch, 1]

float16/bf16

ND

Output, sampled value.

topkToppSamplingType = BATCH_TOPK_MULTINOMIAL_SAMPLING (Batch Random Seed, topk Multinomial Sampling)

Parameter

Dimension

Data Type

Format

Description

probs

[batch, voc_size]

float16/bf16

ND

Input, vocabulary probability.

topk

[batch, 1]

int32

ND

Input top-k, location for top-k truncation. The value of batch must be the same as that of probs.

topp

[batch, 1]

float16/bf16

ND

Input, topp. The value of batch must be the same as that of probs.

sampled_indices

[batch, 1]

int32

ND

Output, sampled idx.

sampled_probs

[batch, 1]

float16/bf16

ND

Output, sampled value.

topkToppSamplingType = BATCH_TOPK_EXPONENTIAL_SAMPLING (Batch Random Seed, topk Exponential Sampling)

Parameter

Dimension

Data Type

Format

Description

probs

[batch, voc_size]

float16/bf16

ND

Input, vocabulary probability.

topk

[batch, 1]

int32

ND

Input top-k, location for top-k truncation. The value of batch must be the same as that of probs.

topp

[batch, 1]

float16/bf16

ND

Input, topp. The value of batch must be the same as that of probs.

exp

[batch, voc_size]

float16/bf16

ND

Input, exponential distribution. The dimension must be the same as that of probs.

sampled_indices

[batch, 1]

int32

ND

Output, sampled idx.

sampled_probs

[batch, 1]

float16/bf16

ND

Output, sampled value.

topkToppSamplingType = BATCH_TOPK_MULTINOMIAL_LOGPROBS_SAMPLING (Batch Random Seed, topk Multinomial Sampling)

Parameter

Dimension

Data Type

Format

Description

probs

[batch, voc_size]

float16/bf16

ND

Input, vocabulary probability.

topk

[batch, 1]

int32

ND

Input top-k, location for top-k truncation. The value of batch must be the same as that of probs.

topp

[batch, 1]

float16/bf16

ND

Input top-p, top-p truncation probability. The value of batch must be the same as that of probs.

rand

[batch, 1]

float

ND

Input. The value range is [0, 1]. For details, see rand() / RAND_MAX.

sampled_indices

[batch, 1]

int32

ND

Output, sampled sequence number.

sampled_probs

[batch, 1]

float16/bf16

ND

Output, sampled value.

sampled_logprobs

[batch, logProbsSize]

float

ND

Output, valid only when topkToppSamplingType is BATCH_TOPK_MULTINOMIAL_LOGPROBS_SAMPLING. The value is the logarithm of the corresponding probability value.

topkToppSamplingType = BATCH_TOPK_EXPONENTIAL_LOGPROBS_SAMPLING (Batch Random Seed, topk Exponential Sampling)

Parameter

Dimension

Data Type

Format

Description

probs

[batch, voc_size]

float16/bf16

ND

Input, vocabulary probability.

topk

[batch, 1]

int32

ND

Input top-k, location for top-k truncation. The value of batch must be the same as that of probs.

topp

[batch, 1]

float16/bf16

ND

Input top-p, top-p truncation probability. The value of batch must be the same as that of probs.

exp

[batch, voc_size]

float16/bf16

ND

Input, exponential distribution. The dimension must be the same as that of probs.

sampled_indices

[batch, 1]

int32

ND

Output, sampled sequence number.

sampled_probs

[batch, 1]

float16/bf16

ND

Output, sampled value.

sampled_logprobs

[batch, logProbsSize]

float

ND

Output, valid only when topkToppSamplingType is BATCH_TOPK_EXPONENTIAL_LOGPROBS_SAMPLING. The value is the logarithm of the corresponding probability value.