ARTICLE · 1080306
TensorRT-LLM 0.5.0 源码之四十八
TensorRT-LLM 0.5.0 源码之四十八点个「赞」+「在看」❤️ 让我们知道这份文字有温暖到你，也是我们持续创作的最大动力！ 推荐 TensorRT-LLM 0.5.0 源码之四十七 SGLang-Omni 通俗解析（六）：新 TTS 车间入驻指南与血泪避坑手册 DALI Audio Resample 把"炼丹炉"直播给你看：小米 MiMo-V2.6-Pro/Flash 强化学习训练全景解读 全双工语音模型究竟为语音AI带来了哪些变革 SGLang-Omni 通俗解析（五）：厂内物流与通信——对讲机只管喊话，物流只管搬货 Moshi：面向实时对话的语音‑文本基础模型 线程别名 Claude Fable 5：自改进智能体——14步循环工程指南 MOSS Transcribe Diarize 技术报告 SGLang-Omni 通俗解析（四）：流水线"四大金刚"与反馈循环黑科技 京东大溶洞 SGLang-Omni 通俗解析（三）：一个请求的奇幻漂流——架构全景 SGLang-Omni 通俗解析（二）：进程命名潜规则与5个高级机制 Qwen3-Omni 技术报告 SGLang-Omni 通俗解析（一）：用"开工厂"的思路理解多模态配置 权限、沙箱与自主智能体 技能：规定智能体应该如何思考 循环工程：设计编码代理系统，而不是每次都手动提示 VoxCPM2 技术报告 Agent 数据生产与训练 什么是循环工程？AI 编码智能体的新范式 从空文件夹到生成图表：Claude Code 实战教程 Claude Code 入门：研究者配置指南 GLM-5.2：面向长时序任务打造 循环工程（Loop Engineering） Agent SFT 标准数据格式 + Loss Mask 完整实现 4D Parallelism TileLang与OpenAI Triton的核心区别 Agent SFT 数据 Claude Code 的上限，就是你的上限 如何用Claude Code提升软件工程工作效率、改善生活 LLM推理优化的核心技术：深入理解KV缓存与分页注意力机制 Qwen3-TTS 技术报告 PagedAttention 如何让AI听懂你的“话外音”？GOAT-SLM模型实现更懂情感的语言交互 FlashAttention与PagedAttention详解：拯救GPU显存，让大模型飞起来的核心技术 
IGptDecoder
class IGptDecoder{public: virtual ~IGptDecoder() = default;virtual void setup(SamplingConfig const& samplingConfig, size_t batchSize)= 0;virtual bool forward(DecodingOutput& output, DecodingInput const& input)= 0;virtual void forwardAsync(DecodingOutput& output, DecodingInput const& input)= 0;static void gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, BufferManager const& manager);static std::unique_ptr<IGptDecoder> create( nvinfer1::DataType dtype, size_t vocabSize, size_t vocabSizePadded, BufferManager::CudaStreamPtr const& stream);};GptDecoder
template <typename T>class GptDecoder : public virtual IGptDecoder{public: using CudaStreamPtr = BufferManager::CudaStreamPtr; GptDecoder(size_t vocabSize, size_t vocabSizePadded, CudaStreamPtr const& stream);void setup(SamplingConfig const& samplingConfig, size_t batchSize) override;bool forward(DecodingOutput& output, DecodingInput const& input) override;void forwardAsync(DecodingOutput& output, DecodingInput const& input) override;private: BufferManager mManager; common::CudaAllocator mAllocator; std::shared_ptr<tensorrt_llm::layers::DynamicDecodeLayer<T>> mDynamicDecodeLayer;};create
inline std::unique_ptr<IGptDecoder> IGptDecoder::create( nvinfer1::DataType dtype, size_t vocabSize, size_t vocabSizePadded, BufferManager::CudaStreamPtr const& stream){ switch (dtype) { case nvinfer1::DataType::kFLOAT: return std::make_unique<GptDecoder<float>>(vocabSize, vocabSizePadded, stream); case nvinfer1::DataType::kHALF: return std::make_unique<GptDecoder<half>>(vocabSize, vocabSizePadded, stream); default: return nullptr; }}GptDecoder
template <typename T>GptDecoder<T>::GptDecoder(size_t vocabSize, size_t vocabSizePadded, CudaStreamPtr const& stream) : mManager{stream} , mAllocator{mManager}{ bool isFreeBufferAfterForward{false}; cudaDeviceProp prop; tc::check_cuda_error(cudaGetDeviceProperties(&prop, 0)); mDynamicDecodeLayer = std::make_shared<tensorrt_llm::layers::DynamicDecodeLayer<T>>( vocabSize, vocabSizePadded, stream->get(), &mAllocator, isFreeBufferAfterForward, &prop);}setup
template <typename T>void GptDecoder<T>::setup(SamplingConfig const& samplingConfig, size_t batchSize){ typename layers::DynamicDecodeLayer<T>::SetupParams setupParams; setupParams.random_seed = samplingConfig.randomSeed; setupParams.repetition_penalty = samplingConfig.repetitionPenalty; setupParams.presence_penalty = samplingConfig.presencePenalty; setupParams.temperature = samplingConfig.temperature; setupParams.min_length = samplingConfig.minLength; // signed to unsigned if (samplingConfig.topK) { auto const& topK = samplingConfig.topK.value(); setupParams.runtime_top_k = std::vector<uint32_t>(std::begin(topK), std::end(topK)); } setupParams.runtime_top_p = samplingConfig.topP; setupParams.top_p_decay = samplingConfig.topPDecay; setupParams.top_p_min = samplingConfig.topPMin; setupParams.top_p_reset_ids = samplingConfig.topPResetIds; setupParams.beam_search_diversity_rate = samplingConfig.beamSearchDiversityRate; setupParams.length_penalty = samplingConfig.lengthPenalty; mDynamicDecodeLayer->setup(batchSize, samplingConfig.beamWidth, setupParams);}forward
template <typename T>bool GptDecoder<T>::forward(DecodingOutput& output, DecodingInput const& input){ TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto forwardParams = prepareInputs<T>(input); auto outputParams = prepareOutputs<T>(output, input.lengths); BufferManager::ITensorPtr finishedSum; std::int32_t* finishedSumHost = nullptr; if (input.sequenceLimitLength && output.finished) { if (output.finishedSum) { finishedSumHost = bufferCast<std::int32_t>(*output.finishedSum); } else { finishedSum = BufferManager::pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); outputParams.finished_sum = tcc::toTllmTensor(*finishedSum); finishedSumHost = bufferCast<std::int32_t>(*finishedSum); } *finishedSumHost = 0; } mDynamicDecodeLayer->forward(outputParams, forwardParams); if (finishedSumHost) { auto const numToFinish = output.finished->getSize(); TLLM_CUDA_CHECK(::cudaStreamSynchronize(mDynamicDecodeLayer->getStream())); return numToFinish == static_cast<std::size_t>(*finishedSumHost); } else { return false; }}template <typename T>typename tl::DynamicDecodeLayer<T>::ForwardParams prepareInputs(DecodingInput const& input){ TLLM_CHECK(input.logits->getDataType() == TRTDataType<T>::value); auto constexpr ite = 0; // no pipeline parallelism typename tl::DynamicDecodeLayer<T>::ForwardParams forwardParams{input.step, ite, input.maxLength, input.batchSize, tcc::toTllmTensor(*input.logits), tcc::toTllmTensor(*input.endIds)}; if (input.cacheIndirection) { forwardParams.src_cache_indirection = tcc::toTllmTensor(*input.cacheIndirection); } if (input.sequenceLimitLength) { forwardParams.sequence_limit_length = tcc::toTllmTensor(*input.sequenceLimitLength); } if (input.embeddingBias) { forwardParams.embedding_bias = tcc::toTllmTensor(*input.embeddingBias); } if (input.lengths) { forwardParams.input_lengths = tcc::toTllmTensor(*input.lengths); } if (input.badWordsList) { forwardParams.bad_words_list = tcc::toTllmTensor(*input.badWordsList); } if (input.stopWordsList) { forwardParams.stop_words_list = tcc::toTllmTensor(*input.stopWordsList); } return forwardParams;}template <typename T>typename tl::DynamicDecodeLayer<T>::OutputParams prepareOutputs( DecodingOutput& output, DecodingInput::TensorPtr const& inputLengths){ TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); typename tl::DynamicDecodeLayer<T>::OutputParams outputParams(tcc::toTllmTensor(*output.ids)); outputParams.newTokens = tcc::toTllmTensor(*output.newTokens); if (output.cumLogProbs) { outputParams.cum_log_probs = tcc::toTllmTensor(*output.cumLogProbs); } if (output.parentIds) { outputParams.parent_ids = tcc::toTllmTensor(*output.parentIds); } if (output.cacheIndirection) { outputParams.tgt_cache_indirection = tcc::toTllmTensor(*output.cacheIndirection); } if (output.finished) { outputParams.finished = tcc::toTllmTensor(*output.finished); } if (output.finishedSum) { outputParams.finished_sum = tcc::toTllmTensor(*output.finishedSum); } if (output.lengths) { outputParams.sequence_length = tcc::toTllmTensor(*output.lengths); } if (output.logProbs) { outputParams.output_log_probs = tcc::toTllmTensor(*output.logProbs); } outputParams.beamHypotheses = std::make_shared<tensorrt_llm::kernels::BeamHypotheses>(); if (output.beamHypotheses.outputIdsTgt) { outputParams.beamHypotheses->output_ids_tgt = bufferCast<int>(*output.beamHypotheses.outputIdsTgt); } if (output.beamHypotheses.sequenceLengthsTgt) { outputParams.beamHypotheses->sequence_lengths_tgt = bufferCast<int>(*output.beamHypotheses.sequenceLengthsTgt); } if (output.beamHypotheses.cumLogProbs) { outputParams.beamHypotheses->cum_log_probs = bufferCast<float>(*output.beamHypotheses.cumLogProbs); } if (output.beamHypotheses.normedScores) { outputParams.beamHypotheses->normed_scores = bufferCast<float>(*output.beamHypotheses.normedScores); } if (output.beamHypotheses.logProbs) { outputParams.beamHypotheses->log_probs = bufferCast<float>(*output.beamHypotheses.logProbs); } if (output.beamHypotheses.minNormedScores) { outputParams.beamHypotheses->min_normed_scores = bufferCast<float>(*output.beamHypotheses.minNormedScores); } if (output.beamHypotheses.numBeams) { outputParams.beamHypotheses->num_beams = bufferCast<int>(*output.beamHypotheses.numBeams); } if (output.beamHypotheses.isDone) { outputParams.beamHypotheses->is_done = bufferCast<bool>(*output.beamHypotheses.isDone); } if (inputLengths) { outputParams.beamHypotheses->input_lengths = bufferCast<int32_t>(*inputLengths); } return outputParams;}forwardAsync
template <typename T>void GptDecoder<T>::forwardAsync(DecodingOutput& output, DecodingInput const& input){ TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto forwardParams = prepareInputs<T>(input); auto outputParams = prepareOutputs<T>(output, input.lengths); mDynamicDecodeLayer->forward(outputParams, forwardParams);}gatherTree
// this should be similar to gatherTree in cpp/tensorrt_llm/thop/gatherTreeOp.cppvoid IGptDecoder::gatherTree(ITensor& finalOutputIds, DecodingOutput const& decodingOutput, DecodingInput const& decodingInput, BufferManager const& manager){ TLLM_LOG_DEBUG("%s start", __PRETTY_FUNCTION__); auto const& finalOutputIdsShape = finalOutputIds.getShape(); auto const& decodingOutputIdsShape = decodingOutput.ids->getShape(); auto const batchSize = finalOutputIdsShape.d[0]; auto const beamWidth = finalOutputIdsShape.d[1]; auto const maxSeqLength = finalOutputIdsShape.d[2]; TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[0] == batchSize, common::fmtstr( "Decoder batch size (%d) does not match final batch size (%d)", decodingOutputIdsShape.d[0], batchSize)); TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[1] == beamWidth, common::fmtstr( "Decoder beam width (%d) does not match final beam width (%d)", decodingOutputIdsShape.d[1], beamWidth)); TLLM_CHECK_WITH_INFO(decodingOutputIdsShape.d[2] <= maxSeqLength, common::fmtstr("Decoder seq length size (%d) is too large for final seq length (%d)", decodingOutputIdsShape.d[2], maxSeqLength)); auto const& stream = manager.getStream(); if (beamWidth > 1) { tensorrt_llm::kernels::invokeInitializeOutput(bufferCast<TokenIdType>(finalOutputIds), bufferCast<TokenIdType>(*decodingInput.endIds), batchSize * beamWidth, maxSeqLength, stream.get()); sync_check_cuda_error(); tensorrt_llm::kernels::BeamHypotheses beamHypotheses; beamHypotheses.sequence_lengths_src = bufferCast<SizeType>(*decodingOutput.lengths); beamHypotheses.parent_ids_src = bufferCast<TokenIdType>(*decodingOutput.parentIds); beamHypotheses.output_ids_src = bufferCast<TokenIdType>(*decodingOutput.ids); beamHypotheses.log_probs_src = nullptr; beamHypotheses.max_seq_len = maxSeqLength; beamHypotheses.length_penalty = 1.0f; beamHypotheses.output_ids_tgt = bufferCast<TokenIdType>(*decodingOutput.beamHypotheses.outputIdsTgt); beamHypotheses.sequence_lengths_tgt = bufferCast<SizeType>(*decodingOutput.beamHypotheses.sequenceLengthsTgt); beamHypotheses.cum_log_probs = bufferCast<float>(*decodingOutput.beamHypotheses.cumLogProbs); beamHypotheses.normed_scores = bufferCast<float>(*decodingOutput.beamHypotheses.normedScores); beamHypotheses.log_probs = bufferCast<float>(*decodingOutput.beamHypotheses.logProbs); beamHypotheses.min_normed_scores = bufferCast<float>(*decodingOutput.beamHypotheses.minNormedScores); beamHypotheses.num_beams = bufferCast<SizeType>(*decodingOutput.beamHypotheses.numBeams); beamHypotheses.is_done = bufferCast<bool>(*decodingOutput.beamHypotheses.isDone); beamHypotheses.input_lengths = bufferCast<SizeType>(*decodingInput.lengths); tensorrt_llm::kernels::invokeInsertUnfinishedPath(beamHypotheses, bufferCast<bool>(*decodingOutput.finished), bufferCast<float>(*decodingOutput.cumLogProbs), batchSize, beamWidth, stream.get()); sync_check_cuda_error(); tensorrt_llm::kernels::invokeFinalize(bufferCast<TokenIdType>(finalOutputIds), bufferCast<SizeType>(*decodingOutput.lengths), bufferCast<float>(*decodingOutput.cumLogProbs), nullptr, // output_logs beamHypotheses.output_ids_tgt, beamHypotheses.sequence_lengths_tgt, beamHypotheses.normed_scores, beamHypotheses.cum_log_probs, beamHypotheses.log_probs, beamHypotheses.num_beams, beamHypotheses.input_lengths, beamWidth, maxSeqLength, batchSize, stream.get()); sync_check_cuda_error(); } else { manager.copy(*decodingOutput.ids, finalOutputIds); sync_check_cuda_error(); }}参考文献
• https://github.com/NVIDIA/TensorRT-LLM/blob/release/0.5.0/cpp/include/tensorrt_llm/runtime/gptDecoder.h