invokeCurandInitialize

__global__ void curandInitialize(curandState_t* state, const int size, const unsigned long long random_seed){    if (threadIdx.x + blockIdx.x * blockDim.x < size)    {        curand_init(random_seed, 0, 0, &state[blockIdx.x * blockDim.x + threadIdx.x]);    }}void invokeCurandInitialize(    curandState_t* state, const size_t batch_size, const unsigned long long random_seed, cudaStream_t stream){    dim3 grid(256);    dim3 block((int) (ceil(batch_size * 1.0 / 256)));    curandInitialize<<<grid, block, 0, stream>>>(state, batch_size, random_seed);}

在CUDA编程中，curand_init 函数用于初始化随机数生成器的状态，是使用cuRAND库在设备端生成随机数的第一步。下面是一个快速总结和详细说明。

💡 核心参数解读

• • 种子 (Seed)：这是随机数序列的“总钥匙”。如果希望两次程序运行产生完全相同的随机数序列（例如，为了复现调试结果），就使用相同的seed。反之，如果想获得不同的随机序列，则应使用不同的seed，例如传入当前时间戳。
• • 序列号 (Sequence)：这个参数在并行计算中至关重要。它为每个线程指定了一个唯一的“子序列标识符”。即使seed相同，只要sequence不同，生成的随机数序列在统计上也是不相关的。这确保了并行线程使用的随机数不会相互干扰。通常，我们将每个线程的全局索引（如 int id = threadIdx.x + blockIdx.x * blockDim.x;）作为其sequence号。
• • 偏移量 (Offset)：你可以把它理解为“快进”按钮。设置offset = n意味着随机数生成器会从序列的第n个位置开始输出随机数。这在需要从序列中特定点开始时非常有用。NVIDIA还提供了skipahead函数，用于在初始化后高效地跳过大量随机数。

📖 基本使用步骤与代码示例

在CUDA核函数中使用curand_init typically遵循以下模式：

1. 1. 声明状态变量：在设备代码中声明一个curandState变量（例如curandStatePhilox4_32_10_t state）。
2. 2. 调用curand_init：为每个线程初始化其独有的状态。需要为它们指定相同的seed、不同的sequence（通常是线程ID），以及可选的offset。
3. 3. 生成随机数：使用初始化后的state，调用具体的分布生成函数，如curand_uniform(&state)来生成[0.0, 1.0)范围内的均匀分布随机浮点数。

下面的代码片段展示了这个过程：

#include <curand_kernel.h>__global__ void myRandomKernel(float *output, unsigned long long seed){    int id = threadIdx.x + blockIdx.x * blockDim.x; // 计算唯一的线程ID    curandStatePhilox4_32_10_t state; // 声明状态变量    // 初始化状态：相同种子，不同序列号（线程ID），偏移量为0    curand_init(seed, id, 0, &state);    // 使用初始化后的状态生成一个[0.0, 1.0)之间的随机数    float randomNum = curand_uniform(&state);     output[id] = randomNum;}

⚡ 性能优化要点

使用curand_init时，有几点性能优化的建议：

• • 初始化开销：curand_init是一个相对耗时的操作。如果内核函数需要多次生成随机数，应避免在循环内重复调用curand_init。正确的做法是在循环外初始化一次，然后在循环内多次使用生成函数。
• • 状态保存：如果一个内核启动后，另一个内核需要继续之前的随机数序列，可以考虑将state状态变量保存到全局内存（Global Memory）中，并在下一个内核中重新加载，而不是重新初始化。这比反复初始化要快得多。
• • 分离内核：为了获得最佳性能，可以考虑将耗时的curand_init初始化过程与调用curand_uniform等生成随机数的过程分离到两个不同的内核函数中。

🌟 支持的其他分布

初始化state后，你不仅可以生成均匀分布的随机数，还可以利用cuRAND Device API生成多种其他分布的随机数，例如：

• • 正态分布：使用 curand_normal(&state)。
• • 对数正态分布：使用 curand_log_normal(&state, mean, stddev)。
• • 泊松分布：使用 curand_poisson(&state, lambda)。

此外，cuRAND还提供了一次性生成2个或4个随机数的向量化函数（如curand_uniform4），这有助于提高生成效率。

invokeCurandBatchInitialize

__global__ void curandBatchInitialize(curandState_t* states, const int size, const unsigned long long* random_seeds){    int idx = threadIdx.x + blockIdx.x * blockDim.x;    if (idx < size)    {        curand_init(random_seeds[idx], 0, 0, &states[idx]);    }}void invokeCurandBatchInitialize(    curandState_t* states, const size_t batch_size, const unsigned long long* random_seeds, cudaStream_t stream){    dim3 grid(256);    dim3 block((int) (ceil(batch_size * 1.0 / 256)));    curandBatchInitialize<<<grid, block, 0, stream>>>(states, batch_size, random_seeds);}

invokeAddBiasEndMask

TopK Sampling steps:

1. 1. invokeAddBiasEndMask
2. 2. invokeAddBiasSoftMax
3. 3. invokeBatchTopKSampling

template <typename T>__global__ void addBiasEndMask(T* logits, const T* bias, const int* end_ids, const bool* finished, const int vocab_size,    const int vocab_size_padded){    int bid = blockIdx.x;    bool finish = finished != nullptr ? finished[bid] : false;    int offset = bid * vocab_size_padded;    const bool IS_FP16 = std::is_same<T, half>::value;    const T MAX_T_VAL = (IS_FP16) ? HALF_FLT_MAX : FLT_MAX;    for (int tid = threadIdx.x; tid < vocab_size_padded; tid += blockDim.x)    {        if (tid >= vocab_size)        {            logits[offset + tid] = -MAX_T_VAL;        }        else if (finish)        {            logits[offset + tid] = (tid == end_ids[bid]) ? MAX_T_VAL : -MAX_T_VAL;        }        else        {            if (bias != nullptr)            {                logits[offset + tid] += bias[tid];            }        }    }}template <typename T>void invokeAddBiasEndMask(T* logits, const T* bias, const int* end_ids, const bool* finished, const int batch_size,    const int vocab_size, const int vocab_size_padded, cudaStream_t stream){    dim3 grid(batch_size);    dim3 block(min(vocab_size_padded, 1024));    /*n is the vocab_size, e.g., 30000, 7000.... vocab_size is usually very big.     */    addBiasEndMask<<<grid, block, 0, stream>>>(logits, bias, end_ids, finished, vocab_size, vocab_size_padded);}template void invokeAddBiasEndMask(float* logits, const float* bias, const int* end_ids, const bool* finished,    const int batch_size, const int vocab_size, const int vocab_size_padded, cudaStream_t stream);template void invokeAddBiasEndMask(half* logits, const half* bias, const int* end_ids, const bool* finished,    const int batch_size, const int vocab_size, const int vocab_size_padded, cudaStream_t stream);

invokeAddBiasSoftMax

template <typename T>__global__ void addBiasSoftMax(    T* logits, const T* bias, const int* end_ids, const bool* finished, const int n_padded, const int n){    int bid = blockIdx.x;    bool finish = (finished != nullptr) ? finished[bid] : false;    int offset = bid * n_padded;    float max_val = -1 * FLT_MAX;    const bool IS_FP16 = std::is_same<T, half>::value;    const T MAX_T_VAL = (IS_FP16) ? HALF_FLT_MAX : FLT_MAX;    __shared__ float s_max_val;    __shared__ float s_sum_val;    for (int tid = threadIdx.x; tid < n_padded; tid += blockDim.x)    {        if (tid < n)        {            if (finish)            {                logits[offset + tid] = (tid == end_ids[bid]) ? MAX_T_VAL : -MAX_T_VAL;            }            else            {                T bias_val = (bias != nullptr) ? bias[tid] : (T) 0.0f;                logits[offset + tid] += bias_val;            }        }        else        {            logits[offset + tid] = -MAX_T_VAL;        }        max_val = max(max_val, (float) logits[offset + tid]);    }    max_val = blockReduceMax<float>((float) max_val);    if (threadIdx.x == 0)    {        s_max_val = max_val;    }    __syncthreads();    float sum_val = 0.0f;    for (int tid = threadIdx.x; tid < n_padded; tid += blockDim.x)    {        logits[offset + tid] = __expf((float) logits[offset + tid] - s_max_val);        sum_val += (float) logits[offset + tid];    }    sum_val = blockReduceSum<float>(sum_val);    if (threadIdx.x == 0)    {        s_sum_val = sum_val;    }    __syncthreads();    for (int tid = threadIdx.x; tid < n_padded; tid += blockDim.x)    {        logits[offset + tid] = ((float) logits[offset + tid] / (s_sum_val + 1e-6f));    }}template <typename T>void invokeAddBiasSoftMax(T* logits, const T* bias, const int* end_ids, const bool* finished, const int m,    const int n_padded, const int n, cudaStream_t stream){    dim3 grid(m);    dim3 block(min(n, 1024));    /*n is the vocab_size, e.g., 30000, 7000.... vocab_size is usually very big.     */    addBiasSoftMax<<<grid, block, 0, stream>>>(logits, bias, end_ids, finished, n_padded, n);}template void invokeAddBiasSoftMax(float* logits, const float* bias, const int* end_ids, const bool* finished,    const int m, const int n_padded, const int n, cudaStream_t stream);template void invokeAddBiasSoftMax(half* logits, const half* bias, const int* end_ids, const bool* finished,    const int m, const int n_padded, const int n, cudaStream_t stream);

invokeBatchTopKSampling

template <typename T>void invokeTopKTopPSampling(void* workspace, size_t& workspace_size, int** output_ids, const T* logits,    int* sequence_lengths, bool* finished_buf, float* cum_log_probs, float* output_log_probs,    curandState_t* curandstate, const int batch_size, const int top_k, const float top_p, const int vocab_size_padded,    const int* end_ids, cudaStream_t stream){    // invokeTopKTopPSampling will be deprecated. Please use invokeTopKSampling    // instead.    invokeTopKSampling(workspace, workspace_size, logits, output_ids, sequence_lengths, finished_buf, cum_log_probs,        output_log_probs, curandstate, top_k, top_p, vocab_size_padded, end_ids, stream, batch_size, nullptr);}template <typename T>void invokeTopKSampling(void* workspace, size_t& workspace_size, const T* log_probs, int** ids, int* sequence_lengths,    bool* finished_buf, float* cum_log_probs, float* output_log_probs, curandState_t* curandstate, const int top_k,    const float top_p, const int vocab_size_padded, const int* end_ids, cudaStream_t stream, const int batch_size,    const bool* skip_decode){    invokeBatchTopKSampling(workspace, workspace_size, log_probs, ids, sequence_lengths, finished_buf, cum_log_probs,        output_log_probs, curandstate, top_k, nullptr, top_p, nullptr, vocab_size_padded, end_ids, stream, batch_size,        skip_decode);}

template <typename T>void invokeBatchTopKSampling(void* workspace, size_t& workspace_size, const T* log_probs, int** ids,    int* sequence_lengths, bool* finished, float* cum_log_probs, float* output_log_probs, curandState_t* curandstate,    const int max_top_k, const int* top_ks, const float top_p, const float* top_ps, const int vocab_size_padded,    const int* end_ids, cudaStream_t stream, const int batch_size, const bool* skip_decode){    TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);    // Not allow an ambiguous inputs top_p and top_ps.    // 只计算topk    assert(top_p == 1.0f || top_ps == nullptr);    const int vocab_size = vocab_size_padded;    const int max_block_per_beam = 8;    int temp_log_probs_buf_size = batch_size * vocab_size;                   // type float    int topk_tmp_ids_buf_size = batch_size * max_top_k * max_block_per_beam; // type int    int topk_tmp_val_buf_size = batch_size * max_top_k * max_block_per_beam; // type float    // prevent memory misaligned address    temp_log_probs_buf_size = (int) (ceil(temp_log_probs_buf_size / 4.)) * 4;    topk_tmp_ids_buf_size = (int) (ceil(topk_tmp_ids_buf_size / 4.)) * 4;    topk_tmp_val_buf_size = (int) (ceil(topk_tmp_val_buf_size / 4.)) * 4;    if (workspace == nullptr)    {        workspace_size = sizeof(T) * temp_log_probs_buf_size + sizeof(int) * topk_tmp_ids_buf_size            + sizeof(T) * topk_tmp_val_buf_size;        return;    }    T* temp_log_probs = (T*) workspace;    int* topk_tmp_id_buf = (int*) (temp_log_probs + temp_log_probs_buf_size);    T* topk_tmp_val_buf = (T*) (topk_tmp_id_buf + topk_tmp_ids_buf_size);    // TODO (bhsueh) need to support case top_k = [2, 17] (use different cases of max_top_k)int log_max_top_k(0);int recursor(max_top_k - 1);    while (recursor >>= 1)        ++log_max_top_k;    switch (log_max_top_k)    {    case 0:    case 1:    case 2:    case 3: // 0 < max_top_k <= 16        CASE_K(16, 128, 128, 8);    case 4: // 16 < max_top_k <= 32        CASE_K(32, 256, 128, 8);    case 5: // 32 < max_top_k <= 64        CASE_K(64, 256, 256, 8);    case 6:    case 7:    case 8:    case 9: // 64 < max_top_k <= 1024        CASE_K(1024, 256, 256, 8);    default: throw std::domain_error(fmtstr("top-k kernel supports 1<=k<=1024 but got k=%d", max_top_k));    }}#define CASE_K(K_MAX, BLOCK_SIZE_1_, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_)                                                  \    topk_stage1<T, BLOCK_SIZE_1_, BLOCKS_PER_BEAM_>                                                                    \        <<<batch_size * BLOCKS_PER_BEAM_, BLOCK_SIZE_1_, 0, stream>>>(log_probs, temp_log_probs, topk_tmp_id_buf,      \            topk_tmp_val_buf, finished, max_top_k, top_ks, vocab_size, end_ids, skip_decode);                          \    topk_stage2_sampling<T, BLOCK_SIZE_2_, BLOCKS_PER_BEAM_>                                                           \        <<<batch_size, BLOCK_SIZE_2_, K_MAX * sizeof(int) + K_MAX * sizeof(float), stream>>>(topk_tmp_id_buf,          \            topk_tmp_val_buf, ids, sequence_lengths, finished, cum_log_probs, output_log_probs, max_top_k, top_ks,     \            top_p, top_ps, curandstate, end_ids, vocab_size, skip_decode);                                             \    break;#undef CASE_K

这段代码用于根据 max_top_k（即 Top-K 算法中的 K 值）的大小，动态选择一组优化的配置参数（如 GPU 核函数中的线程块大小、网格大小等）。其核心思想是将 K 值按 2 的幂次进行分档，为不同范围的 K 值匹配不同的计算资源，以优化性能。下面逐部分解释。

🔢 1. 计算对数尺度（log_max_top_k）

int log_max_top_k(0);int recursor(max_top_k - 1);while (recursor >>= 1)    ++log_max_top_k;

• • 目的：计算 max_top_k - 1 的二进制位宽（即最高有效位的位置），相当于求 log2(max_top_k) 的整数部分。这是一种高效的近似方法，用于将 K 值映射到对数尺度。
• • 过程：
• • 初始化 recursor 为 max_top_k - 1（避免 K=1 时直接右移得到 0 的问题）。
• • 循环通过右移操作（>>= 1）不断将 recursor 除以 2，直到其为 0。每次右移，log_max_top_k 加 1。
• • 示例：若 max_top_k = 16，则 recursor = 15（二进制 1111），需右移 4 次才变为 0，因此 log_max_top_k = 4。

⚙️ 2. Switch-Case 分档配置

switch (log_max_top_k) {    case 0: case 1: case 2:    case 3:  // 对应 0 < max_top_k ≤ 16        CASE_K(16, 128, 128, 8);    case 4:  // 16 < max_top_k ≤ 32        CASE_K(32, 256, 128, 8);    ...}

• • 目的：根据 log_max_top_k 的值，跳转到对应的配置档位。每个档位对应一个 K 值范围（见注释），并通过宏 CASE_K 设置具体参数。
• • Fall-through 设计：
• • case 0、1、2 后没有 break，会直接执行到 case 3 的代码。这意味着 K ≤ 16 时统一使用 CASE_K(16, ...) 的配置。这种设计是为了合并小 K 值的处理，减少代码冗余。
• • 类似地，case 6–9 也合并到 case 9 的处理（K ≤ 1024）。
• • 参数含义（推测）：CASE_K 宏的参数可能表示：
• • 第一个参数（如 16、32）：当前档位支持的最大 K 值。
• • 后续参数（如 128、256）：GPU 核函数的配置（如线程块大小、网格大小、共享内存大小等）。例如，CASE_K(1024, 256, 256, 8) 可能表示当 K 较大时，使用 256 个线程的块大小和 8 个块等。

⚠️ 3. 错误处理（default case）

default:    throw std::domain_error(fmtstr("top-k kernel supports 1<=k<=1024 but got k=%d", max_top_k));

• • 当 log_max_top_k > 9（即 max_top_k > 1024）时，抛出异常。这是因为代码只预设了 K ≤ 1024 的配置，更大的 K 可能需其他优化方式。

💡 代码设计思想

• • 性能优化：Top-K 算法（如基于堆或快速排序的变种）的性能受数据规模影响。通过分档配置，可以为不同规模的 K 匹配最佳计算资源（如线程数），避免资源浪费。
• • 可扩展性：如需支持 K > 1024，可添加类似 case 10: CASE_K(2048, 512, 256, 8); 的档位。

🔍 实际应用场景

这类代码常见于 GPU 加速的 Top-K 算法实现（如深度学习中的采样操作）。通过动态选择核函数参数，确保在不同 K 值下均能高效利用并行计算资源。

topk_stage1

计算topk

template <typename T, int BLOCK_SIZE_, int BLOCKS_PER_BEAM_>__global__ void topk_stage1(const T* __restrict log_probs, T* tmp_log_probs, int* topk_tmp_id_buf, T* topk_tmp_val_buf,    const bool* finished, const int max_top_k, const int* top_ks, const int vocab_size, const int* end_ids,    const bool* skip_decode){    // log_probs，其实是 softmax 或 logits 的结果。取决于 cum_log_probs和output_log_probs 的值。    typedef cub::BlockReduce<TopK_2<T>, BLOCK_SIZE_> BlockReduce;    __shared__ typename BlockReduce::TempStorage temp_storage;    const int tid = threadIdx.x;    const int bid = blockIdx.x;    const int batch_id = bid / BLOCKS_PER_BEAM_; // row id for log_probs    if (skip_decode != nullptr && skip_decode[batch_id])    {        return;    }    const int block_lane = bid % BLOCKS_PER_BEAM_;                    // block id for a beam    const int k = (top_ks != nullptr) ? top_ks[batch_id] : max_top_k; // batch_id = batch index    const int tmp_log_buf_index = batch_id * vocab_size; // [B, V]    const int tmp_topk_buf_index = batch_id * BLOCKS_PER_BEAM_ * max_top_k + block_lane * k; // [Batch, Beam(BLOCKS_PER_BEAM_), k]    TopK_2<T> partial;    const bool IS_FP16 = std::is_same<T, half>::value;    const T MAX_T_VAL = (IS_FP16) ? HALF_FLT_MAX : FLT_MAX;    if (finished != nullptr && finished[batch_id] == true)    {   //finished        if (tid < k)        {            const int index = tmp_topk_buf_index + tid;            if (block_lane == 0 && tid == 0)            {   // top1                const int end_id = end_ids[batch_id];                // output                topk_tmp_id_buf[index] = tmp_log_buf_index + end_id;                topk_tmp_val_buf[index] = log_probs[tmp_log_buf_index + end_id];            }            else            {   // other k-1                // output                topk_tmp_id_buf[index] = -1;                topk_tmp_val_buf[index] = -MAX_T_VAL;            }        }        return;    }     // copy beam 的 score， [1, V]    for (int elem_id = tid + block_lane * BLOCK_SIZE_; elem_id < vocab_size; elem_id += BLOCK_SIZE_ * BLOCKS_PER_BEAM_)    {        int index = elem_id + tmp_log_buf_index;        tmp_log_probs[index] = log_probs[index]; // tmp output, [B, V]. 用于排序用    }    for (int ite = 0; ite < k; ite++)    {        partial.init();#pragma unroll        for (int elem_id = tid + block_lane * BLOCK_SIZE_; elem_id < vocab_size;             elem_id += BLOCK_SIZE_ * BLOCKS_PER_BEAM_)        {            // topk(elem_id, ..., elem_id+= BLOCK_SIZE_ * BLOCKS_PER_BEAM_)            int index = elem_id + tmp_log_buf_index;            partial.insert(tmp_log_probs[index], index);        }        TopK_2<T> total = BlockReduce(temp_storage).Reduce(partial, reduce_topk_op_2<T>);        if (tid == 0)        {            const int index = tmp_topk_buf_index + ite; // Top-i            topk_tmp_id_buf[index] = total.p; // [Batch, Beam(BLOCKS_PER_BEAM_), max_top_k]            topk_tmp_val_buf[index] = total.u; // [Batch, Beam(BLOCKS_PER_BEAM_), max_top_k]            if (total.p >= 0)            {                tmp_log_probs[total.p] = -MAX_T_VAL; // 已经top{N-1}的score设置最小，便于计算topN            }        }        __syncthreads();    }}

topk_stage2_sampling

1. 1. 计算topp+sampling
2. 2. sampling 用的是 curand_uniform (0.0, 1.0]

template <typename T, int BLOCK_SIZE_, int BLOCKS_PER_BEAM_>__global__ void topk_stage2_sampling(const int* __restrict topk_tmp_id_buf, T* topk_tmp_val_buf, int** ids,    int* sequence_lengths, bool* finished, float* cum_log_probs, float* output_log_probs, const int max_top_k,    const int* top_ks, const float top_p, const float* top_ps, curandState_t* curandstate, const int* end_ids,    const int vocab_size, const bool* skip_decode){    const bool IS_FP16 = std::is_same<T, half>::value;    const T MAX_T_VAL = (IS_FP16) ? HALF_FLT_MAX : FLT_MAX;    const int tid = threadIdx.x;    const int batch_id = blockIdx.x;    if (skip_decode != nullptr && skip_decode[batch_id])    {        return;    }    const int k = (top_ks != nullptr) ? top_ks[batch_id] : max_top_k;    const float prob_threshold = (top_ps != nullptr) ? top_ps[batch_id] : top_p;    const int size = k * BLOCKS_PER_BEAM_;    const int stride = max_top_k * BLOCKS_PER_BEAM_;    typedef cub::BlockReduce<TopK_2<float>, BLOCK_SIZE_> BlockReduce;    __shared__ typename BlockReduce::TempStorage temp_storage;    extern __shared__ char array[];  // K_MAX * sizeof(int) +K_MAX* sizeof(float)    __shared__ float rand_num;    __shared__ float s_sum;    __shared__ float s_max;    // topk_tmp_val_buf shape [Batch, Beam, max_top_k]    T* s_val = topk_tmp_val_buf + batch_id * stride;    int* s_id = reinterpret_cast<int*>(array); // shm index    if (tid == 0)    {        s_sum = 0.0f;    }    TopK_2<float> partial;    if (finished != nullptr && finished[batch_id] == true)    {        ids[batch_id][sequence_lengths[batch_id]] = end_ids[batch_id];        return;    }    float* s_val2 = reinterpret_cast<float*>(s_id + k); // shm score    for (int ite = 0; ite < k; ite++)    {        partial.init();#pragma unroll        for (int i = tid; i < size; i += BLOCK_SIZE_) // all BLOCKS_PER_BEAM_ beams, [BLOCKs_PER_BEAM, k]        {            partial.insert((float) s_val[i], i);        }        TopK_2<float> total = BlockReduce(temp_storage).Reduce(partial, reduce_topk_op_2<float>);        if (tid == 0)        {            if (ite == 0)            {                s_max = total.u;            }            s_id[ite] = total.p;            s_val[total.p] = -MAX_T_VAL;            // when cum_log_probs are computed, topk_tmp_val_buf (logits_buf_) are            // already pre-processed by softmax_kernel            if (cum_log_probs == nullptr && output_log_probs == nullptr)            {                total.u = __expf(total.u - s_max);            }            s_val2[ite] = total.u; // prob            s_sum += total.u; // cum topk probs        }        __syncthreads();    }    if (tid == 0)    {        // (0.0, 1.0]        rand_num = (float) curand_uniform(curandstate + blockIdx.x) * prob_threshold * s_sum;  // 随机数        for (int i = 0; i < k; i++)        {            float exp_logit = s_val2[i];            rand_num = rand_num - exp_logit; // remainder cum prob            if (rand_num <= 0.0f || i == k - 1)            {                // s_id[i] in range [0, BLOCK_PSER_BEAM_ *k)                ids[batch_id][sequence_lengths[batch_id]] = topk_tmp_id_buf[batch_id * stride + s_id[i]] % vocab_size;                if (cum_log_probs != nullptr || output_log_probs != nullptr)                {                    float log_prob = logf(exp_logit); // compute log prob                    if (cum_log_probs != nullptr)                    {                        // 不是 log_prob 的累计。log里面没有除以sum_logit                        cum_log_probs[batch_id] += log_prob;                    }                    if (output_log_probs != nullptr)                    {                        // 'output_log_probs' is the probability induced by the top-k                        // sampling. We normalize the probability 'exp_logit' of the                        // selected token by the probability 's_sum' of a set of top-k                        // tokens, meaning the log_prob is the probability of the selected                        // token, conditioned on the event that it is selected, i.e.,                        //   log_prob = log P(i | i is in top-k) = log(exp_logit / s_sum).                        output_log_probs[batch_id] = log_prob - logf(s_sum);                    }                }                break;            }        }        if (sequence_lengths != nullptr && finished != nullptr)        {            int seqlen = sequence_lengths[batch_id];            finished[batch_id] = ids[batch_id][seqlen] == end_ids[batch_id];            if (!finished[batch_id])            {                sequence_lengths[batch_id] = seqlen + 1;            }        }    }}

samplingTopX

// topk，当k==0，且p!=0时，跳过计算。skip_decode[i] = k == 0;// topp, 当k>0, 跳过计算.skip_decode[i] = k > 0;

参考文献

• • https://github.com/NVIDIA/TensorRT-LLM/blob/v0.5.0/cpp/tensorrt_llm/kernels/samplingTopKKernels.h