baseLayer

class BaseLayer{public:    BaseLayer(cudaStream_t stream, tensorrt_llm::common::IAllocator* allocator, bool is_free_buffer_after_forward,        cudaDeviceProp* cuda_device_prop = nullptr)        : stream_(stream)        , allocator_(allocator)        , cuda_device_prop_(cuda_device_prop)        , is_free_buffer_after_forward_(is_free_buffer_after_forward){};    virtual ~BaseLayer() = default;virtual cudaStream_t getStream(){        return stream_;    }virtual void setStream(cudaStream_t stream){        stream_ = stream;    }protected:    // device environments    cudaStream_t stream_;    tensorrt_llm::common::IAllocator* allocator_;    cudaDeviceProp* cuda_device_prop_ = nullptr;    bool is_free_buffer_after_forward_;    bool is_allocate_buffer_ = false; // TODO to be deprecated};

DecodingSetupParams

namespace tc = tensorrt_llm::common;class DecodingSetupParams{public:    std::optional<std::vector<float>> temperature;       // [1] or [batch_size] on cpu    std::optional<std::vector<std::int32_t>> min_length; // [1] or [batch_size] on cpu    // repetition_penalty and presence_penalty are mutually exclusive.    std::optional<std::vector<float>> repetition_penalty; // [1] or [batch_size] on cpu    std::optional<std::vector<float>> presence_penalty;   // [1] or [batch_size] on cpu};

DecodingParams

class DecodingParams{public:    DecodingParams(int step, int ite, tc::Tensor logits, tc::Tensor end_ids)        : step{step}        , ite{ite}        , logits{std::move(logits)}        , end_ids{std::move(end_ids)}    {    }    // mandatory parameters    int step;    int ite; // iteration. ite 用于标识当前正在处理的是这个子批次中的第几个独立序列或实例。    tc::Tensor logits;  // [local_batch_size, beam_width, vocab_size_padded]    tc::Tensor end_ids; // [local_batch_size]};

DecodingOutputParams

class DecodingOutputParams{public:explicit DecodingOutputParams(tc::Tensor outputIds)        : output_ids{std::move(outputIds)}    {    }    // mandatory parameters    tc::Tensor output_ids; // [max_seq_len, batch_size]    // optional parameters    std::optional<tc::Tensor> finished;        // [batch_size * beam_width], optional    std::optional<tc::Tensor> sequence_length; // [batch_size * beam_width], optional    std::optional<tc::Tensor> cum_log_probs;   // [batch_size * beam_width], necessary in beam search    std::optional<tc::Tensor>        output_log_probs;                 // [request_ouptut_length, batch_size * beam_width], must be float*, optional    std::optional<tc::Tensor> parent_ids; // [max_seq_len, batch_size * beam_width], necessary in beam search    tc::Tensor output_ids_ptr;            // [batch_size] int* (2-d array), each int* has [beam_width, max_seq_len]};

BaseSamplingLayer

template <typename T>class BaseSamplingLayer : public BaseLayer{public:    BaseSamplingLayer(size_t vocab_size, size_t vocab_size_padded, cudaStream_t stream,        tensorrt_llm::common::IAllocator* allocator, bool is_free_buffer_after_forward,        cudaDeviceProp* cuda_device_prop);    BaseSamplingLayer(BaseSamplingLayer const& sampling_layer);    ~BaseSamplingLayer() override;    class SetupParams : public DecodingSetupParams    {    public:        std::optional<std::vector<std::uint32_t>> runtime_top_k;    // [1] or [batch_size] on cpu        std::optional<std::vector<float>> runtime_top_p;            // [1] or [batch_size] on cpu        std::optional<std::vector<unsigned long long>> random_seed; // [1] or [batch_size] on cpu    };    class ForwardParams : public DecodingParams    {    public:        ForwardParams(int step, int ite, tc::Tensor logits, tc::Tensor end_ids, int max_seq_len)            : DecodingParams{step, ite, std::move(logits), std::move(end_ids)}            , max_seq_len{max_seq_len}        {        }        // mandatory parameters        int max_seq_len;        // optional parameters        std::optional<tc::Tensor> embedding_bias; // [vocab_size_padded]        std::optional<tc::Tensor> input_lengths;  // [local_batch_size * beam_width]    };void forward(DecodingOutputParams& outputs, ForwardParams const& params);protected:    size_t vocab_size_;    size_t vocab_size_padded_;    size_t sampling_workspace_size_;    void* sampling_workspace_ = nullptr;    // device    curandState_t* curandstate_buf_ = nullptr; // [B,]    unsigned long long* random_seeds_buf_ = nullptr; // [B,]    float* temperature_buf_ = nullptr; // [B,]    float* repetition_penalty_buf_ = nullptr; // [B,]    int32_t* min_lengths_buf_ = nullptr; // [B,]    bool* skip_decode_buf_ = nullptr; // [B,]    T* runtime_logits_buf_ = nullptr;  // [B,V]      // host    std::vector<float> mTemperature;    std::vector<float> mRepetitionPenalty;    std::vector<int32_t> mMinLengths;    bool* skip_decode_ = nullptr;  // [B,], host    bool skip_any_ = false;    tensorrt_llm::kernels::RepetitionPenaltyType repetition_penalty_type_        = tensorrt_llm::kernels::RepetitionPenaltyType::None;virtual void runSampling(DecodingOutputParams& outputs, DecodingParams const& params)= 0;virtual void freeBuffer();void setupBase(size_t batch_size, SetupParams const& setupParams);private:void allocateBuffer(size_t batch_size);bool isValidBatchSize(size_t batch_size);};

template <typename T>BaseSamplingLayer<T>::BaseSamplingLayer(size_t vocab_size, size_t vocab_size_padded, cudaStream_t stream,    IAllocator* allocator, bool is_free_buffer_after_forward, cudaDeviceProp* cuda_device_prop)    : BaseLayer(stream, allocator, is_free_buffer_after_forward, cuda_device_prop)    , vocab_size_(vocab_size)    , vocab_size_padded_(vocab_size_padded){}template <typename T>BaseSamplingLayer<T>::BaseSamplingLayer(BaseSamplingLayer const& sampling_layer)    : BaseLayer(sampling_layer)    , vocab_size_(sampling_layer.vocab_size_)    , vocab_size_padded_(sampling_layer.vocab_size_padded_)    , sampling_workspace_size_(sampling_layer.sampling_workspace_size_){}template <typename T>BaseSamplingLayer<T>::~BaseSamplingLayer(){}

allocateBuffer

template <typename T>void BaseSamplingLayer<T>::allocateBuffer(size_t batch_size){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    curandstate_buf_ = allocator_->reMalloc(curandstate_buf_, sizeof(curandState_t) * batch_size, false); // [B,]    random_seeds_buf_ = allocator_->reMalloc(random_seeds_buf_, sizeof(unsigned long long) * batch_size, false); // [B,]    temperature_buf_ = allocator_->reMalloc(temperature_buf_, sizeof(float) * batch_size, false); // [B,]    repetition_penalty_buf_ = allocator_->reMalloc(repetition_penalty_buf_, sizeof(float) * batch_size, false); // [B,]    min_lengths_buf_ = allocator_->reMalloc(min_lengths_buf_, sizeof(int) * batch_size, false); // [B,]    runtime_logits_buf_ = allocator_->reMalloc(runtime_logits_buf_, sizeof(T) * batch_size * vocab_size_padded_, false); // [B,V]    skip_decode_buf_ = allocator_->reMalloc(skip_decode_buf_, sizeof(bool) * batch_size, false); // [B,]    // host buffers.    skip_decode_ = (bool*) std::realloc(skip_decode_, sizeof(bool) * batch_size);  // [B,]    TLLM_CHECK(skip_decode_ != nullptr);    is_allocate_buffer_ = true;}

template <typename T>void BaseSamplingLayer<T>::freeBuffer(){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    if (is_allocate_buffer_)    {        allocator_->free((void**) (&curandstate_buf_));        allocator_->free((void**) (&random_seeds_buf_));        allocator_->free((void**) (&temperature_buf_));        allocator_->free((void**) (&repetition_penalty_buf_));        allocator_->free((void**) (&min_lengths_buf_));        allocator_->free((void**) (&runtime_logits_buf_));        allocator_->free((void**) (&skip_decode_buf_));        std::free(skip_decode_);        is_allocate_buffer_ = false;    }}

setupBase

default random_seed = 0

template <typename T>void BaseSamplingLayer<T>::setupBase(const size_t batch_size, SetupParams const& setupParams){    TLLM_LOG_DEBUG(__PRETTY_FUNCTION__);    allocateBuffer(batch_size);    // If runtime argument has single random seed, using this random seed to    // initialize the random table of all sentences. If the argument has    // [batch_size] random seeds, initializing the random table by different    // random seeds respectively. If no random seed, initialize the random table    // of all sentences by 0 directly.    // default seed 0.    if (setupParams.random_seed)    {        if (setupParams.random_seed->size() == 1)        {            invokeCurandInitialize(curandstate_buf_, batch_size, setupParams.random_seed->front(), stream_);            sync_check_cuda_error();        }        else        {            TLLM_CHECK_WITH_INFO(setupParams.random_seed->size() == batch_size, "Random seed vector size mismatch.");            cudaAutoCpy(random_seeds_buf_, setupParams.random_seed->data(), batch_size, stream_);            invokeCurandBatchInitialize(curandstate_buf_, batch_size, random_seeds_buf_, stream_);            sync_check_cuda_error();        }    }    else    {        // Initialize curand states using the default seed 0.        invokeCurandInitialize(curandstate_buf_, batch_size, 0, stream_);    }    // Setup penalties.    auto fillBuffers        = [this, &batch_size](auto const& optParam, auto const defaultValue, auto& hostBuffer, auto& deviceBuffer)    {        hostBuffer.resize(batch_size);        if (!optParam)        {            std::fill(std::begin(hostBuffer), std::end(hostBuffer), defaultValue);        }        else if (optParam->size() == 1)        {            std::fill(std::begin(hostBuffer), std::end(hostBuffer), optParam->front());        }        else        {            TLLM_CHECK_WITH_INFO(optParam->size() == batch_size, "Argument vector size mismatch.");            std::copy(optParam->begin(), optParam->end(), std::begin(hostBuffer));        }        cudaAutoCpy(deviceBuffer, hostBuffer.data(), batch_size, stream_);    };      // default temperature 1.0      // defualt min_length 1    fillBuffers(setupParams.temperature, 1.0f, mTemperature, temperature_buf_);    fillBuffers(setupParams.min_length, 1, mMinLengths, min_lengths_buf_);    // default repetition_penalty 0.0    if ((setupParams.repetition_penalty) || (setupParams.presence_penalty))    {        TLLM_CHECK_WITH_INFO(!((setupParams.repetition_penalty) && (setupParams.presence_penalty)),            "Found ambiguous parameters repetition_penalty and presence_penalty "            "which are mutually exclusive. "            "Please provide one of repetition_penalty or presence_penalty.");        repetition_penalty_type_ = (setupParams.repetition_penalty) ? RepetitionPenaltyType::Multiplicative                                                                    : RepetitionPenaltyType::Additive;        auto const& repetition_penalty = (repetition_penalty_type_ == RepetitionPenaltyType::Multiplicative)            ? setupParams.repetition_penalty            : setupParams.presence_penalty;        // repetition_penalty is not empty, so default value 0.0f has no effect        fillBuffers(repetition_penalty, 0.0f, mRepetitionPenalty, repetition_penalty_buf_);    }    else    {        repetition_penalty_type_ = RepetitionPenaltyType::None;    }}

penaltyTypes

enum class RepetitionPenaltyType{    Additive,       // the presence penalty    Multiplicative, // the repetition penalty    None            // No repetition penalty.};inline float getDefaultPenaltyValue(RepetitionPenaltyType penalty_type){    switch (penalty_type)    {    case RepetitionPenaltyType::Additive: return 0.0f;    case RepetitionPenaltyType::Multiplicative: return 1.0f;    default: break;    }    return 0.0f;}

forward

template class BaseSamplingLayer<float>; // fp32template class BaseSamplingLayer<half>;  // fp16, no bf16

template <typename T>void BaseSamplingLayer<T>::forward(DecodingOutputParams& outputs, ForwardParams 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;    auto const step = params.step;    auto* const input_lengths = params.input_lengths ? params.input_lengths->template getPtr<const int>() : nullptr;    auto* logits = params.logits.template getPtr<T>();#define ALL_OF(p_, sz_, dt_, v_) (std::all_of(p_, p_ + sz_, [&](dt_ b) { return b == v_; }))    bool* skip_decode = skip_decode_ + ite * local_batch_size;    if (ALL_OF(skip_decode, local_batch_size, bool, true))    {        // No sample in the current batch to do TopX sampling.        return;    }    skip_any_ = std::any_of(skip_decode, skip_decode + local_batch_size, [](bool b) { return b; });    if (skip_any_)    {        // A TopX Sampling layer directly changes the logit values. In case of        // skip_any==true, meaning topk and topp layers will run simultaneously for        // a batch in the same step. We copy the logits to an internal buffer, not        // affecting the other sampling layers.        TLLM_CHECK(params.logits.size() == local_batch_size * vocab_size_padded_);        cudaD2Dcpy(runtime_logits_buf_, logits, params.logits.size());        logits = runtime_logits_buf_;    }    // 1. logits embedding bias or temperature    auto* embedding_bias = params.embedding_bias ? params.embedding_bias->template getPtr<T const>() : nullptr;    if (embedding_bias != nullptr        || !ALL_OF(std::begin(mTemperature) + ite * local_batch_size, local_batch_size, float, 1.0f))    {        invokeBatchApplyTemperaturePenalty(logits, embedding_bias, temperature_buf_ + ite * local_batch_size,            local_batch_size, vocab_size_, vocab_size_padded_, stream_);    }    sync_check_cuda_error();        // 2. repeation_penalty    if (step > 1 && repetition_penalty_type_ != RepetitionPenaltyType::None)    {        float default_value = getDefaultPenaltyValue(repetition_penalty_type_);        if (!ALL_OF(std::begin(mRepetitionPenalty) + ite * local_batch_size, local_batch_size, float, default_value))        {            auto* const input_lengths                = params.input_lengths ? params.input_lengths->template getPtr<const int>() : nullptr;            invokeBatchApplyRepetitionPenalty(logits, repetition_penalty_buf_ + ite * local_batch_size,                outputs.output_ids_ptr.template getPtr<const int*>(), outputs.sequence_length->getPtr<const int>(),                batch_size, local_batch_size, vocab_size_padded_, input_lengths, repetition_penalty_type_,                params.max_seq_len, stream_);            sync_check_cuda_error();        }    }     // 3. minLengthPenalty    auto* end_ids = params.end_ids.template getPtr<const int>();    invokeMinLengthPenalty(logits, min_lengths_buf_ + ite * local_batch_size, end_ids,        outputs.sequence_length->getPtr<const int>(), input_lengths, local_batch_size, vocab_size_padded_, stream_);    sync_check_cuda_error();#undef ALL_OF     // 4. TopX sampling    runSampling(outputs, params);    if (is_free_buffer_after_forward_)    {        freeBuffer();    }    sync_check_cuda_error();    TLLM_LOG_DEBUG("%s stop", __PRETTY_FUNCTION__);}

参考文献

• • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/layers/CMakeLists.txt
• • https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/tensorrt_llm/layers/baseSamplingLayer.h