TopKSamplingLayer

template <typename T>class TopKSamplingLayer : public BaseSamplingLayer<T>{public:    using Base = BaseSamplingLayer<T>;    using SetupParams = typename Base::SetupParams;    TopKSamplingLayer(size_t vocab_size, size_t vocab_size_padded, cudaStream_t stream,        tensorrt_llm::common::IAllocator* allocator, bool is_free_buffer_after_forward);    TopKSamplingLayer(TopKSamplingLayer<T> const& top_k_sampling_layer);    ~TopKSamplingLayer();void setup(size_t batch_size, SetupParams const& setupParams);protected:void runSampling(DecodingOutputParams& outputs, DecodingParams const& params) override;void freeBuffer() override;    uint32_t runtime_max_top_k_ = 1;    uint32_t* runtime_top_k_buf_ = nullptr;    float* runtime_top_p_buf_ = nullptr;    using Base::vocab_size_;    using Base::vocab_size_padded_;    using Base::sampling_workspace_size_;    using Base::sampling_workspace_;    using Base::curandstate_buf_;    using Base::random_seeds_buf_;    using Base::skip_decode_buf_;    using Base::skip_decode_;    using Base::skip_any_;    using Base::runtime_logits_buf_;    using Base::stream_;    using Base::allocator_;    using Base::is_allocate_buffer_;private:void allocateBuffer(size_t batch_size, std::vector<uint32_t> const& top_k);};

template <typename T>TopKSamplingLayer<T>::TopKSamplingLayer(size_t vocab_size, size_t vocab_size_padded, cudaStream_t stream,    IAllocator* allocator, bool is_free_buffer_after_forward)    : BaseSamplingLayer<T>(vocab_size, vocab_size_padded, stream, allocator, is_free_buffer_after_forward, nullptr){}template <typename T>TopKSamplingLayer<T>::TopKSamplingLayer(TopKSamplingLayer<T> const& top_k_sampling_layer)    : BaseSamplingLayer<T>(top_k_sampling_layer){}template <typename T>TopKSamplingLayer<T>::~TopKSamplingLayer(){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    freeBuffer();}template class TopKSamplingLayer<float>;template class TopKSamplingLayer<half>;

allocateBuffer

top_k为0，底层会使用greedy_decode， 即top_k=1

template <typename T>void TopKSamplingLayer<T>::allocateBuffer(size_t const batch_size, std::vector<uint32_t> const& top_k){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    uint32_t max_top_k = (top_k.size() > 0) ? *std::max_element(std::begin(top_k), std::end(top_k)) : 1;    if (max_top_k == 0)    {        // for safety. TopKSamplingLayer handles a case of top_k=0 and top_p=0 as        // a greedy decode, i.e. top_k=1, although such case has max_top_k=0.        max_top_k = 1;    }    invokeTopKSampling<T>(nullptr, sampling_workspace_size_, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,        nullptr, max_top_k, 1.0f, vocab_size_padded_, nullptr, stream_, batch_size, skip_decode_buf_);    sampling_workspace_ = allocator_->reMalloc(sampling_workspace_, sampling_workspace_size_, false);    runtime_top_k_buf_ = allocator_->reMalloc(runtime_top_k_buf_, sizeof(uint32_t) * batch_size, false); // [B,]    runtime_top_p_buf_ = allocator_->reMalloc(runtime_top_p_buf_, sizeof(float) * batch_size, false); // [B,]    is_allocate_buffer_ = true;}

template <typename T>void TopKSamplingLayer<T>::freeBuffer(){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    if (is_allocate_buffer_)    {        allocator_->free((void**) (&sampling_workspace_));        allocator_->free((void**) (&runtime_top_k_buf_));        allocator_->free((void**) (&runtime_top_p_buf_));    }    BaseSamplingLayer<T>::freeBuffer();    is_allocate_buffer_ = false;}

setup

1. 1. topk=0 and top=0.0, 即greedy decode.
• • 等价于 topk=1, topp=0.0
• • 等价于 topk=1, topp=1.0
2. 2. 最大 topk = 1024
3. 3. 0.0 <= topp <= 1.0f
4. 4. topk=0, skip_decode=True

template <typename T>void TopKSamplingLayer<T>::setup(size_t const batch_size, SetupParams const& setupParams){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    BaseSamplingLayer<T>::setupBase(batch_size, setupParams);    uint32_t const default_top_k = 0;    auto const runtime_top_k = setupParams.runtime_top_k.value_or(std::vector<uint32_t>{default_top_k});    auto const runtime_top_p = setupParams.runtime_top_p.value_or(std::vector<float>{});    allocateBuffer(batch_size, runtime_top_k);    size_t const runtime_top_k_size = runtime_top_k.size();    size_t const runtime_top_p_size = runtime_top_p.size();    uint32_t const top_k = *std::max_element(std::begin(runtime_top_k), std::end(runtime_top_k));    float const top_p = (runtime_top_p_size == 0) ? 0.0f : runtime_top_p.front();    if (runtime_top_k_size > 1)    {        TLLM_CHECK_WITH_INFO(runtime_top_k.size() == batch_size,            fmtstr(                "runtime_top_k.size() (%lu) == batch_size (%lu) is not satisfied!", runtime_top_k.size(), batch_size));        cudaAutoCpy(runtime_top_k_buf_, runtime_top_k.data(), batch_size, stream_);    }    if (runtime_top_p_size > 1)    {        TLLM_CHECK_WITH_INFO(runtime_top_p.size() == batch_size,            fmtstr(                "runtime_top_p.size() (%lu) == batch_size (%lu) is not satisfied!", runtime_top_p.size(), batch_size));        cudaAutoCpy(runtime_top_p_buf_, runtime_top_p.data(), batch_size, stream_);    }    dim3 block(std::min((int) batch_size, 256));    dim3 grid(divUp((int) batch_size, (int) block.x));    // support top_k up to 1024.    setup_topk_runtime_args<1024><<<grid, block, 0, stream_>>>(batch_size, top_k, runtime_top_k_buf_,        runtime_top_k_size, top_p, runtime_top_p_buf_, runtime_top_p_size, skip_decode_buf_);    cudaAutoCpy(skip_decode_, skip_decode_buf_, batch_size, stream_);    std::vector<uint32_t> runtime_top_ks(batch_size);    cudaAutoCpy(runtime_top_ks.data(), runtime_top_k_buf_, batch_size, stream_);    runtime_max_top_k_ = *std::max_element(std::begin(runtime_top_ks), std::end(runtime_top_ks));}

设置 runtime_top_k_buf_, runtime_top_p_buf_, skip_decode_buf_参数

template <uint32_t TOP_K_MAX>__global__ void setup_topk_runtime_args(int batch_size, uint32_t top_k, uint32_t* top_ks, int top_ks_size, float top_p,    float* top_ps, int top_ps_size, bool* skip_decode){    int index = blockIdx.x * blockDim.x + threadIdx.x;    for (int i = index; i < batch_size; i += gridDim.x * blockDim.x)    {        uint32_t k = top_ks_size > 1 ? top_ks[i] : top_k;        float p = top_ps_size > 1 ? top_ps[i] : top_p;        if (k == 0 && p == 0.0f)        {            // TensorRT-LLM's topp implementation does not support topp = 0.0f, but it            // equivalent to greedy search. So, we set the topk = 1 as an alternative            // solution.            k = 1;        }        if (k > 0 && p == 0.0f)        {            // for compatibility <= TensorRT-LLM5.0.            // This case corresponds to the old topk sampling, which is equivalent to            // the old topk_topp sampling with topp=1.0f. TopKSamplingLayer and            // TopKTopPSamplingLayer are now merged by TopKSamplingLayer. Thus, we            // replace the case topk>0 and topp=0.0f by topk>0 and topp=1.0f for the            // compatibility.            p = 1.0f;        }        // Clip k value. A topk sampling kernel supports up to TOP_K_MAX=64.        top_ks[i] = k > TOP_K_MAX ? TOP_K_MAX : k;        if (k > TOP_K_MAX)        {            printf(                "[WARNING] topk (%d) is larger than max supported number (%d) for "                "token %d"                " clip to max supported number %d. \n",                k, TOP_K_MAX, i, top_ks[i]);        }        // Clip p value if it is out of range. range = [0.0, 1.0].        top_ps[i] = p < 0.0f ? 0.0f : (p > 1.0f ? 1.0f : p);        if (p < 0.0f || p > 1.0f)        {            printf(                "[WARNING] topp (%f) is out of range ([0.0, 1.0f]) for token %d"                " clip to closest number %f.\n",                p, i, top_ps[i]);        }        skip_decode[i] = k == 0;    }}

runSampling

template <typename T>void TopKSamplingLayer<T>::runSampling(DecodingOutputParams& outputs, DecodingParams const& params){    TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__);    auto const batch_size = outputs.output_ids_ptr.shape[0];    auto const local_batch_size = params.logits.shape[0];    auto const ite = params.ite;    // in case of skip any, the logit value is already copied and processed.    auto* logits = !skip_any_ ? params.logits.template getPtr<T>() : runtime_logits_buf_;    auto* end_ids = params.end_ids.template getPtr<const int>();    bool* finished = (outputs.finished) ? outputs.finished->template getPtr<bool>() : nullptr;    invokeAddBiasEndMask(        logits, (T*) (nullptr), end_ids, finished, local_batch_size, vocab_size_, vocab_size_padded_, stream_);    sync_check_cuda_error();    float* cum_log_probs = (outputs.cum_log_probs) ? outputs.cum_log_probs->template getPtr<float>() : nullptr;    float* output_log_probs = (outputs.output_log_probs) ? outputs.output_log_probs->template getPtr<float>() : nullptr;    if (cum_log_probs != nullptr || output_log_probs != nullptr)    {        invokeAddBiasSoftMax(            logits, (T*) (nullptr), end_ids, finished, local_batch_size, vocab_size_padded_, vocab_size_, stream_);        sync_check_cuda_error();    }    int* sequence_length = (outputs.sequence_length) ? outputs.sequence_length->template getPtr<int>() : nullptr;    invokeBatchTopKSampling(sampling_workspace_, sampling_workspace_size_, logits,        outputs.output_ids_ptr.template getPtr<int*>(), sequence_length, finished, cum_log_probs, output_log_probs,        curandstate_buf_ + ite * local_batch_size,        (int) runtime_max_top_k_, // useless because runtime_top_k_buf_ is never                                  // nullptr. Keep for legacy.        (int*) (runtime_top_k_buf_ + ite * local_batch_size),        1.0f,                     // useless because runtime_top_p_buf_ is never nullptr. Keep for                                  // legacy.        runtime_top_p_buf_ + ite * local_batch_size, vocab_size_padded_, end_ids, stream_, local_batch_size,        skip_decode_buf_ + ite * local_batch_size);    sync_check_cuda_error();}

TopPSamplingLayer

template <typename T>class TopPSamplingLayer : public BaseSamplingLayer<T>{public:    using Base = BaseSamplingLayer<T>;    class SetupParams : public Base::SetupParams    {    public:        std::optional<std::vector<float>> top_p_decay;            // [batch_size], must between [0, 1]        std::optional<std::vector<float>> top_p_min;              // [batch_size], must between [0, 1]        std::optional<std::vector<std::int32_t>> top_p_reset_ids; // [batch_size]    };    TopPSamplingLayer(std::size_t vocab_size, std::size_t vocab_size_padded, cudaStream_t stream,        tensorrt_llm::common::IAllocator* allocator, bool is_free_buffer_after_forward,        cudaDeviceProp* cuda_device_prop);    TopPSamplingLayer(TopPSamplingLayer<T> const& top_p_sampling_layer);    ~TopPSamplingLayer();void setup(std::size_t batch_size, SetupParams const& setupParams);protected:void runSampling(DecodingOutputParams& outputs, DecodingParams const& params) override;void freeBuffer() override;    std::uint32_t* runtime_top_k_buf_ = nullptr;    float* runtime_top_p_buf_ = nullptr;    float runtime_max_top_p_;    float* initial_top_p_buf_ = nullptr;    float* top_p_decay_buf_ = nullptr;    float* top_p_min_buf_ = nullptr;    std::int32_t* top_p_reset_ids_buf_ = nullptr;    std::int32_t* topp_id_vals_buf_ = nullptr;    std::int32_t* topp_offset_buf_ = nullptr;    std::int32_t* begin_topp_offset_buf_ = nullptr;    std::size_t cub_temp_storage_size_;    using Base::vocab_size_;    using Base::vocab_size_padded_;    using Base::sampling_workspace_size_;    using Base::sampling_workspace_;    using Base::curandstate_buf_;    using Base::random_seeds_buf_;    using Base::skip_decode_buf_;    using Base::skip_decode_;    using Base::skip_any_;    using Base::runtime_logits_buf_;    using Base::stream_;    using Base::allocator_;    using Base::is_allocate_buffer_;    using Base::cuda_device_prop_;private:void allocateBuffer(std::size_t batch_size, std::vector<float> const& top_k);};

template <typename T>TopPSamplingLayer<T>::TopPSamplingLayer(std::size_t vocab_size, std::size_t vocab_size_padded, cudaStream_t stream,    IAllocator* allocator, bool is_free_buffer_after_forward, cudaDeviceProp* cuda_device_prop)    : BaseSamplingLayer<T>(        vocab_size, vocab_size_padded, stream, allocator, is_free_buffer_after_forward, cuda_device_prop){}template <typename T>TopPSamplingLayer<T>::TopPSamplingLayer(TopPSamplingLayer<T> const& top_p_sampling_layer)    : BaseSamplingLayer<T>(top_p_sampling_layer){}template <typename T>TopPSamplingLayer<T>::~TopPSamplingLayer(){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    freeBuffer();}template class TopPSamplingLayer<float>;template class TopPSamplingLayer<half>;

allocateBuffer

template <typename T>void TopPSamplingLayer<T>::allocateBuffer(std::size_t batch_size, std::vector<float> const& top_p){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    float const max_top_p = (top_p.size() > 0) ? *std::max_element(std::begin(top_p), std::end(top_p)) : 0.0f;    invokeTopPSampling<T>(nullptr, // workspace        sampling_workspace_size_, cub_temp_storage_size_,        nullptr,                   // output_ids        nullptr,                   // sequence_length        nullptr,                   // finished_buffer        nullptr,                   // cum_log_probs        nullptr,                   // output_log_probs        nullptr,                   // log_probs        topp_id_vals_buf_, topp_offset_buf_, begin_topp_offset_buf_, curandstate_buf_, batch_size, vocab_size_padded_,        nullptr, max_top_p, stream_, cuda_device_prop_, skip_decode_buf_);    sampling_workspace_ = allocator_->reMalloc(sampling_workspace_, sampling_workspace_size_, true);    runtime_top_k_buf_ = allocator_->reMalloc(runtime_top_k_buf_, sizeof(std::uint32_t) * batch_size, false);    runtime_top_p_buf_ = allocator_->reMalloc(runtime_top_p_buf_, sizeof(float) * batch_size, false);    initial_top_p_buf_ = allocator_->reMalloc(initial_top_p_buf_, sizeof(float) * batch_size, false);    top_p_decay_buf_ = allocator_->reMalloc(top_p_decay_buf_, sizeof(float) * batch_size, false);    top_p_min_buf_ = allocator_->reMalloc(top_p_min_buf_, sizeof(float) * batch_size, false);    top_p_reset_ids_buf_ = allocator_->reMalloc(top_p_reset_ids_buf_, sizeof(std::int32_t) * batch_size, false);    topp_id_vals_buf_        = allocator_->reMalloc(topp_id_vals_buf_, sizeof(std::int32_t) * batch_size * vocab_size_padded_, false);    topp_offset_buf_ = allocator_->reMalloc(topp_offset_buf_, sizeof(std::int32_t) * (batch_size + 1), false);    begin_topp_offset_buf_        = allocator_->reMalloc(begin_topp_offset_buf_, sizeof(std::int32_t) * (batch_size + 1), false);    is_allocate_buffer_ = true;}

template <typename T>void TopPSamplingLayer<T>::freeBuffer(){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    if (is_allocate_buffer_)    {        allocator_->free((void**) (&sampling_workspace_));        allocator_->free((void**) (&topp_id_vals_buf_));        allocator_->free((void**) (&topp_offset_buf_));        allocator_->free((void**) (&begin_topp_offset_buf_));        allocator_->free((void**) (&runtime_top_k_buf_));        allocator_->free((void**) (&runtime_top_p_buf_));        allocator_->free((void**) (&initial_top_p_buf_));        allocator_->free((void**) (&top_p_decay_buf_));        allocator_->free((void**) (&top_p_min_buf_));        allocator_->free((void**) (&top_p_reset_ids_buf_));    }    BaseSamplingLayer<T>::freeBuffer();    is_allocate_buffer_ = false;}

setup

1. 1. topk=0 and topp=0.0, greedy decode topk=1
2. 2. topk > 0 , skip_decode = True
3. 3. default top_p_decay = 1.0
4. 4. default top_p_min = 0.5f

template <typename T>void TopPSamplingLayer<T>::setup(std::size_t const batch_size, SetupParams const& setupParams){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    BaseSamplingLayer<T>::setupBase(batch_size, setupParams);    std::uint32_t const default_top_k = 0;    auto const runtime_top_k = setupParams.runtime_top_k.value_or(std::vector<uint32_t>{default_top_k});    auto const runtime_top_p = setupParams.runtime_top_p.value_or(std::vector<float>{});    allocateBuffer(batch_size, runtime_top_p);    std::size_t const runtime_top_k_size = runtime_top_k.size();    std::size_t const runtime_top_p_size = runtime_top_p.size();    if (runtime_top_p_size == 0)    {        std::fill_n(skip_decode_, batch_size, true);        return;    }    std::uint32_t const top_k = runtime_top_k.at(0);    float const top_p = runtime_top_p.at(0);    if (runtime_top_k_size > 1)    {        TLLM_CHECK_WITH_INFO(runtime_top_k.size() == batch_size,            fmtstr(                "runtime_top_k.size() (%lu) == batch_size (%lu) is not satisfied!", runtime_top_k.size(), batch_size));        cudaAutoCpy(runtime_top_k_buf_, runtime_top_k.data(), batch_size, stream_);    }    if (runtime_top_p_size > 1)    {        TLLM_CHECK_WITH_INFO(runtime_top_p.size() == batch_size,            fmtstr(                "runtime_top_p.size() (%lu) == batch_size (%lu) is not satisfied!", runtime_top_p.size(), batch_size));        cudaAutoCpy(runtime_top_p_buf_, runtime_top_p.data(), batch_size, stream_);    }    auto fillBuffers = [this, &batch_size](std::string name, auto const& vector, auto& deviceBuffer)    {        TLLM_CHECK_WITH_INFO(vector.size() == batch_size,            fmtstr("%s.size() (%lu) == batch_size (%lu) is not satisfied!", name.c_str(), vector.size(), batch_size));        cudaAutoCpy(deviceBuffer, vector.data(), batch_size, stream_);    };    float const defaultTopPDecay{1.0f};    fillBuffers("top_p_decay", setupParams.top_p_decay.value_or(std::vector<float>(batch_size, defaultTopPDecay)),        top_p_decay_buf_);    float const defaultTopPMin{1e-6f}; // prevent topp becoming 0.0    fillBuffers(        "top_p_min", setupParams.top_p_min.value_or(std::vector<float>(batch_size, defaultTopPMin)), top_p_min_buf_);    std::int32_t const defaultTopPResetId{-1};    fillBuffers("top_p_reset_ids",        setupParams.top_p_reset_ids.value_or(std::vector<std::int32_t>(batch_size, defaultTopPResetId)),        top_p_reset_ids_buf_);    dim3 block(std::min((int) batch_size, 256));    dim3 grid(divUp((int) batch_size, (int) block.x));    set_topp_runtime_args<<<grid, block, 0, stream_>>>(batch_size, top_k, runtime_top_k_buf_, runtime_top_k_size, top_p,        runtime_top_p_buf_, runtime_top_p_size, skip_decode_buf_, initial_top_p_buf_, top_p_decay_buf_, top_p_min_buf_);    sync_check_cuda_error();    cudaAutoCpy(skip_decode_, skip_decode_buf_, batch_size, stream_);    std::vector<float> runtime_top_ps(batch_size);    cudaAutoCpy(runtime_top_ps.data(), runtime_top_p_buf_, batch_size, stream_);    runtime_max_top_p_ = *std::max_element(std::begin(runtime_top_ps), std::end(runtime_top_ps));}

static __global__ void set_topp_runtime_args(int batch_size, std::uint32_t top_k, std::uint32_t* top_ks,    int top_ks_size, float top_p, float* top_ps, int top_ps_size, bool* skip_decode, float* initial_top_p_buf,    float* top_p_decay_buf, float* top_p_min_buf){    /**     * @brief Setup the runtime arguments for topp, broadcasting top_p to top_ps              and top_k to top_ks, verifying value ranges of top_p_decay/top_p_min.     *     * \param batch_size     * \param top_k                 first top_k     * \param top_ks                [batch_size]     * \param top_ks_size     * \param top_p                 first top_p     * \param top_ps                [batch_size]     * \param top_ps_size     * \param skip_decode           [batch_size]     * \param initial_top_p_buf     [batch_size]     * \param top_p_decay_buf       [batch_size]     * \param top_p_min_buf         [batch_size]     *     */    int index = blockIdx.x * blockDim.x + threadIdx.x;    for (int i = index; i < batch_size; i += gridDim.x * blockDim.x)    {        std::uint32_t k = top_ks_size > 1 ? top_ks[i] : top_k;        float p = top_ps_size > 1 ? top_ps[i] : top_p;        if (k == 0 && p == 0.0f)        {            // TensorRT-LLM's topp implementation does not support topp = 0.0f, but it            // equivalent to greedy search. So, we set the topk = 1 as an alternative            // solution.            k = 1;        }        top_ks[i] = k;        // Clip p value if it is out of range. range = [0.0, 1.0].        top_ps[i] = p < 0.0f ? 0.0f : (p > 1.0f ? 1.0f : p);        if (p < 0.0f || p > 1.0f)        {            printf(                "[WARNING] topp (%f) is out of range ([0.0, 1.0f]) for token %d"                " clip to closest number %f.\n",                p, i, top_ps[i]);        }        skip_decode[i] = k > 0;  // Warning        initial_top_p_buf[i] = top_ps[i];        if (top_p_decay_buf[i] > 1.0f || top_p_decay_buf[i] <= 0.0f)        {            printf(                "[WARNING] top_p_decay_buf (%f) is out of range ([0.0, 1.0f]) for "                "token %d,"                " change to 1.0f.\n",                top_p_decay_buf[i], i);            top_p_decay_buf[i] = 1.0f;        }        if (top_p_min_buf[i] > 1.0f || top_p_min_buf[i] <= 0.0f)        {            printf(                "[WARNING] top_p_min_buf (%f) is out of range ([0.0, 1.0f]) for "                "token %d,"                " change to 0.5f.\n",                top_p_min_buf[i], i);            top_p_min_buf[i] = 0.5f;        }    }}

runSampling

template <typename T>void TopPSamplingLayer<T>::runSampling(DecodingOutputParams& outputs, DecodingParams const& params){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    auto const batch_size = outputs.output_ids_ptr.shape[0];    auto const local_batch_size = params.logits.shape[0];    auto const ite = params.ite;    // in case of skip any, the logit value is already copied and processed.    auto* logits = !skip_any_ ? params.logits.template getPtr<T>() : runtime_logits_buf_;    auto* end_ids = params.end_ids.template getPtr<const int>();    invokeTopPInitialize(        topp_id_vals_buf_, topp_offset_buf_, begin_topp_offset_buf_, local_batch_size, vocab_size_padded_, stream_);    sync_check_cuda_error();    bool* finished = (outputs.finished) ? outputs.finished->template getPtr<bool>() : nullptr;    invokeAddBiasSoftMax(        logits, (T*) (nullptr), end_ids, finished, local_batch_size, vocab_size_padded_, vocab_size_, stream_);    sync_check_cuda_error();    float* cum_log_probs = (outputs.cum_log_probs) ? outputs.cum_log_probs->template getPtr<float>() : nullptr;    float* output_log_probs = (outputs.output_log_probs) ? outputs.output_log_probs->template getPtr<float>() : nullptr;    int* sequence_length = (outputs.sequence_length) ? outputs.sequence_length->template getPtr<int>() : nullptr;    invokeBatchTopPSampling<T>(sampling_workspace_, sampling_workspace_size_, cub_temp_storage_size_,        outputs.output_ids_ptr.template getPtr<int*>(), sequence_length, finished, cum_log_probs, output_log_probs,        logits, topp_id_vals_buf_, topp_offset_buf_, begin_topp_offset_buf_, curandstate_buf_ + ite * local_batch_size,        local_batch_size, vocab_size_padded_, end_ids, runtime_max_top_p_, runtime_top_p_buf_ + ite * local_batch_size,        stream_, cuda_device_prop_, skip_decode_buf_ + ite * local_batch_size);    sync_check_cuda_error();    invokeComputeToppDecay(runtime_top_p_buf_ + ite * local_batch_size, initial_top_p_buf_ + ite * local_batch_size,        outputs.output_ids_ptr.template getPtr<const int*>(), top_p_decay_buf_ + ite * local_batch_size,        top_p_min_buf_ + ite * local_batch_size, top_p_reset_ids_buf_ + ite * local_batch_size, sequence_length,        local_batch_size, stream_);    sync_check_cuda_error();}

参考文献

• • https://github.com/NVIDIA/TensorRT-LLM/blob/v0.5.0/cpp/tensorrt_llm/kernels/samplingPenaltyKernels.h
• • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/layers/topPSamplingLayer.h