夜雨聆风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 bf16template <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
