运行时

bufferManager.h

namespace tensorrt_llm

namespace runtime

class BufferManager

#include <bufferManager.h>

一个用于管理主机和设备内存的辅助类。

公共类型

using IBufferPtr = IBuffer::UniquePtr

using ITensorPtr = ITensor::UniquePtr

using CudaStreamPtr = std::shared_ptr<CudaStream>

using CudaMemPoolPtr = std::shared_ptr<CudaMemPool>

公共函数

explicit BufferManager(CudaStreamPtr stream, bool trimPool = false)

构造一个 BufferManager。

参数:

cudaStream – [in] 用于 GPU 上所有操作（分配、释放、复制等）的 cuda 流。

inline ~BufferManager()

析构函数。

IBufferPtr gpu(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

) const

使用 cudaMallocAsync 在 GPU 上分配给定大小的 IBuffer。

ITensorPtr gpu(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

) const

使用 cudaMallocAsync 在 GPU 上分配给定维度的 ITensor。

IBufferPtr allocate(

MemoryType memoryType,

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

) const

分配给定大小和内存类型的 IBuffer。

ITensorPtr allocate(

MemoryType memoryType,

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

) const

分配给定维度和内存类型的 ITensor。

inline IBufferPtr emptyBuffer(

MemoryType memoryType,

nvinfer1::DataType type = kBYTE_TYPE,

) const

创建给定内存类型的空 IBuffer。以后可以调整其大小。

inline ITensorPtr emptyTensor(

MemoryType memoryType,

nvinfer1::DataType type = kBYTE_TYPE,

) const

创建给定内存类型的空 ITensor。以后可以重塑其维度。

void setMem(IBuffer &buffer, int32_t value) const

将给定 buffer 的内容设置为指定值。

void setZero(IBuffer &buffer) const

将给定 buffer 的内容设置为零。

void copy(void const *src, IBuffer &dst, MemoryType srcType) const

将 src 复制到 dst。

void copy(IBuffer const &src, void *dst, MemoryType dstType) const

将 src 复制到 dst。

inline void copy(void const *src, IBuffer &dst) const

将 src 复制到 dst。

inline void copy(IBuffer const &src, void *dst) const

将 src 复制到 dst。

void copy(IBuffer const &src, IBuffer &dst) const

将 src 复制到 dst。

IBufferPtr copyFrom(IBuffer const &src, MemoryType memoryType) const

将 src 复制到一个新的 IBuffer 中，新 buffer 的内存类型可能不同。

ITensorPtr copyFrom(ITensor const &src, MemoryType memoryType) const

将 src 复制到一个新的 ITensor 中，新 tensor 的内存类型可能不同。

template<typename T>
inline IBufferPtr copyFrom(

std::vector<T> const &src,

MemoryType memoryType,

) const

将 src 复制到一个新的 IBuffer 中，新 buffer 的内存类型可能不同。

template<typename T>
inline ITensorPtr copyFrom(

T *src,

nvinfer1::Dims dims,

MemoryType memoryType,

) const

将 src 复制到一个新的 ITensor 中，新 tensor 的内存类型可能不同。

template<typename T>
inline ITensorPtr copyFrom(

std::vector<T> const &src,

nvinfer1::Dims dims,

MemoryType memoryType,

) const

将 src 复制到一个新的 ITensor 中，新 tensor 的内存类型可能不同。

CudaStream const &getStream() const

获取底层的 cuda 流。

std::size_t memoryPoolReserved() const

内存池保留的当前内存大小。

std::size_t memoryPoolUsed() const

内存池使用的当前内存大小。

std::size_t memoryPoolFree() const

内存池中当前可用的内存大小。

void memoryPoolTrimTo(std::size_t size)

尝试将内存池保留的内存大小修剪到 size 字节。此操作会与流隐式同步。

公共静态函数

static IBufferPtr gpuSync(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 GPU 上使用 cudaMalloc 分配给定大小的 IBuffer。

static ITensorPtr gpuSync(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 GPU 上使用 cudaMalloc 分配给定维度的 ITensor。

static IBufferPtr cpu(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 上分配给定大小的 IBuffer。

static ITensorPtr cpu(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 上分配给定维度的 ITensor。

static IBufferPtr pinned(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 上分配给定大小的固定 IBuffer。

static ITensorPtr pinned(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 上分配给定维度的固定 ITensor。

static IBufferPtr pinnedPool(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 的默认内存池中分配给定大小的固定 IBuffer。

static ITensorPtr pinnedPool(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 CPU 的默认内存池中分配给定维度的固定 ITensor。

static IBufferPtr managed(

std::size_t size,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 UVM 中分配给定大小的 IBuffer。

static ITensorPtr managed(

nvinfer1::Dims dims,

nvinfer1::DataType type = kBYTE_TYPE,

)

在 UVM 中分配给定维度的 ITensor。

static ITensorPtr ipcNvls(

std::set<int> ranks,

nvinfer1::Dims dims,

nvinfer1::DataType type,

)

为 NVLS 分配给定维度的 ITensor。

公共静态属性

static auto constexpr kBYTE_TYPE = nvinfer1::DataType::kUINT8

私有成员

CudaStreamPtr mStream

CudaMemPoolPtr mPool

bool const mTrimPool

友元

friend class ::BufferManagerTest

common.h

宏定义

FMT_DIM

namespace tensorrt_llm

namespace runtime

类型定义

using SizeType32 = std::int32_t

using SizeType64 = std::int64_t

using TokenIdType = std::int32_t

using LoraTaskIdType = std::uint64_t

using TokenExtraIdType = std::uint64_t

using VecTokenExtraIds = std::vector<TokenExtraIdType>

using VecUniqueTokens = std::vector<UniqueToken>

template<typename T>
using StringPtrMap = std::unordered_map<std::string, std::shared_ptr<T>>

枚举

enum class RequestType : std::int32_t

值

enumerator kCONTEXT

enumerator kGENERATION

struct UniqueToken

公共函数

inline bool operator==(UniqueToken const &other) const noexcept

公共成员

TokenIdType tokenId

TokenExtraIdType tokenExtraId

cudaEvent.h

namespace tensorrt_llm

namespace runtime

class CudaEvent

公共类型

using pointer = cudaEvent_t

公共函数

inline explicit CudaEvent(unsigned int flags = cudaEventDisableTiming)

创建一个新的 CUDA 事件。该事件将在析构函数中被销毁。

参数:

flags – 事件创建标志。默认情况下，事件计时被禁用。

inline explicit CudaEvent(pointer event, bool ownsEvent = true)

将现有 CUDA 事件传递给此对象。

参数:

• event – 要传递给此对象的事件。

• ownsEvent – 此对象是否拥有该事件并在析构函数中销毁它。

inline pointer get() const

返回与此对象关联的事件。

inline void synchronize() const

同步事件。

私有类型

using element_type = std::remove_pointer_t<pointer>

using EventPtr = std::unique_ptr<element_type, Deleter>

私有成员

EventPtr mEvent

class Deleter

公共函数

inline explicit Deleter(bool ownsEvent)

inline explicit Deleter()

inline constexpr void operator()(pointer event) const

私有成员

bool mOwnsEvent

cudaStream.h

namespace tensorrt_llm

namespace runtime

class CudaStream

公共函数

inline explicit CudaStream(

unsigned int flags = cudaStreamNonBlocking,

int priority = 0,

)

在当前设备上创建一个新的 CUDA 流。此流将在析构函数中销毁。

参数:

• flags – 流创建标志。请参阅 cudaStreamCreateWithFlags 了解可以传递的有效标志列表。

• priority – 流的优先级。数字越小代表优先级越高。请参阅 cudaDeviceGetStreamPriorityRange 了解可以传递的有意义的流优先级的更多信息。

inline explicit CudaStream(

cudaStream_t stream,

int device,

bool ownsStream = true,

)

将现有 CUDA 流传递给此对象。

参数:

• stream – 要传递给此对象的流。

• device – 创建此流的设备。

• ownsStream – 此对象是否拥有该流并在析构函数中销毁它。

inline explicit CudaStream(cudaStream_t stream)

使用现有 CUDA 流或通过传递 nullptr 使用默认流进行构造。

inline int getDevice() const

返回创建此流的设备。

inline cudaStream_t get() const

返回与此对象关联的流。

inline void synchronize() const

同步此流。

inline void record(CudaEvent::pointer event) const

在此流上记录一个事件。

inline void record(CudaEvent const &event) const

在此流上记录一个事件。

inline void wait(CudaEvent::pointer event) const

等待一个事件。

inline void wait(CudaEvent const &event) const

等待一个事件。

私有类型

using StreamPtr = std::unique_ptr<std::remove_pointer_t<cudaStream_t>, Deleter>

私有成员

StreamPtr mStream

int mDevice = {-1}

友元

friend class CudaStreamBindings

class Deleter

公共函数

inline explicit Deleter(bool ownsStream)

inline explicit Deleter()

inline constexpr void operator()(cudaStream_t stream) const

私有成员

bool mOwnsStream

decodingInput.h

namespace tensorrt_llm

namespace runtime

class DecodingInput

#include <decodingInput.h>

表示解码器的输入。

此输入类型被假定为不可变。它表示解码器最初接收到的内容，并始终可以这样引用。

公共类型

using TensorConstPtr = ITensor::SharedConstPtr

using TensorPtr = ITensor::SharedPtr

公共函数

inline DecodingInput(

SizeType32 maxLength,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 batchSize,

TensorConstPtr logits,

TensorPtr endIds,

TensorConstPtr batchSlots,

)

公共成员

SizeType32 step

强制参数 我们正在进行的解码步骤的索引。仅在 Python 运行时中使用。

SizeType32 maxLength

要解码的最大 token 数量。

SizeType32 maxAttentionWindow

解码时要考虑的最大注意力窗口长度。

SizeType32 sinkTokenLength

用作注意力汇聚点的 token 数量，https://arxiv.org/html/2309.17453v3。

SizeType32 batchSize

批次中的样本数量。

std::vector<SizeType32> beamWidths

每个请求的束宽度，[batchSize]。

SizeType32 maxStopWordsLen

`stopWordsLens` 张量中的最大值。

SizeType32 maxBadWordsLen

`badWordsLens` 张量中的最大值。

TensorConstPtr logits

模型前向计算的输出，是词汇表上的概率分布 [batchSize, beamWidth, vocabSizePadded]，位于 gpu 上。

std::optional<std::vector<TensorConstPtr>> logitsVec

对 logits 的另一种视图，[batchSize][beamWidth, vocabSizePadded]，位于 gpu 上。

TensorConstPtr endIds

结束 ID，[batchSize * beamWidth]，位于 gpu 上。

TensorConstPtr batchSlots

线性批次 ID 到序列槽位的地址映射，[batchSize]，位于 pinned 内存，类型为 int32_t。

TensorConstPtr finishReasons

可选参数 当前迭代的完成状态（如果为 true，则跳过请求的解码步骤），[batchSize, beamWidth]，位于 gpu 上。

TensorConstPtr sequenceLimitLength

批次中每个序列的最大序列长度，[batchSize]，位于 gpu 上。

TensorConstPtr embeddingBias

TensorConstPtr lengths

std::vector<TensorPtr> badWordsLists

TensorConstPtr badWordsPtrs

TensorConstPtr badWordsLens

std::vector<TensorPtr> stopWordsLists

TensorConstPtr stopWordsPtrs

TensorConstPtr stopWordsLens

TensorConstPtr noRepeatNgramSize

TensorPtr cacheIndirection

用于束搜索的 KV 缓存索引的参数，GPU 上形状为 [batchSize, beamWidth, maxSeqLen]

std::optional<MedusaInputs> medusaInputs

std::optional<ExplicitDraftTokensInputs> explicitDraftTokensInputs

std::optional<LookaheadInputs> lookaheadInputs

std::optional<ExternalDraftTokensInputs> externalDraftTokensInputs

std::optional<EagleInputs> eagleInputs

struct EagleInputs

公共函数

inline EagleInputs(

TensorConstPtr nextDraftTokens,

TensorConstPtr nextDraftLens,

TensorConstPtr nextDraftPaths,

TensorConstPtr lastDraftTokens,

TensorConstPtr lastDraftLens,

TensorConstPtr lastDraftPaths,

TensorConstPtr acceptedTokens,

TensorConstPtr acceptedLens,

TensorConstPtr acceptedPathIds,

TensorConstPtr chunkedContextNextTokens,

TensorConstPtr seqSlots,

)

公共成员

TensorConstPtr nextDraftTokens

TensorConstPtr nextDraftLens

TensorConstPtr nextDraftPaths

TensorConstPtr lastDraftTokens

TensorConstPtr lastDraftLens

TensorConstPtr lastDraftPaths

TensorConstPtr acceptedTokens

TensorConstPtr acceptedLens

TensorConstPtr acceptedPathIds

TensorConstPtr chunkedContextNextTokens

TensorConstPtr seqSlots

class ExplicitDraftTokensInputs

公共成员

TensorConstPtr nextDraftTokens

TensorConstPtr nextFlatTokens

TensorConstPtr nextDraftIndices

TensorConstPtr nextDraftProbs

TensorConstPtr lastDraftTokens

TensorConstPtr lastDraftIndices

TensorConstPtr masks

TensorConstPtr packedPositionIds

TensorConstPtr bestPathLengths

TensorConstPtr bestPathIndices

TensorConstPtr nextGenerationLengths

TensorConstPtr lastPositionIdsBase

TensorConstPtr lastGenerationLengths

TensorConstPtr maxGenLengthDevice

TensorConstPtr seqSlots

class ExternalDraftTokensInputs

公共成员

TensorPtr draftLogits

TensorPtr draftProbs

TensorPtr targetProbs

TensorPtr numDraftTokens

TensorPtr numDraftTokensHost

TensorPtr draftTokenIds

TensorPtr useDraftLogits

TensorPtr useDraftLogitsHost

SizeType32 step

float constantThreshold

bool useRandomAcceptanceThreshold

struct LookaheadInputs

公共成员

TensorPtr tokensPerStep

class MedusaInputs

公共成员

TensorConstPtr medusaPaths

[batchSize, maxTokensPerStep, maxMedusaHeads + 1]，在 GPU 上

TensorConstPtr medusaTreeIds

[batchSize, maxTokensPerStep]，在 GPU 上

std::vector<std::vector<TensorPtr>> medusaLogits

[batchSize][maxAcceptedDraftTokensPerStep][maxDraftTokens + 1, vocabSizePadded]，在 GPU 上

TensorPtr medusaCurTokensPerStep

[batchSize]，在 GPU 上

TensorConstPtr medusaTargetTokensPerStep

[batchSize]，在 GPU 上

decodingOutput.h

namespace tensorrt_llm

namespace batch_manager

namespace runtime

class DecodingOutput

公共类型

using TensorPtr = ITensor::SharedPtr

公共函数

inline explicit DecodingOutput(TensorPtr ids, TensorPtr gatheredIds)

公共成员

TensorPtr ids

TensorPtr gatheredIds

TensorPtr newTokensSteps

TensorPtr newTokens

std::vector<TensorPtr> newTokensVec

TensorPtr finishReasons

TensorPtr finishedSum

TensorPtr logProbs

TensorPtr cumLogProbs

TensorPtr parentIds

TensorPtr lengths

TensorPtr cacheIndirection

TensorPtr logProbsTiled

BeamHypotheses beamHypotheses

std::optional<SpeculativeDecodingOutputs> speculativeDecodingOutputs

std::optional<ExplicitDraftTokensBuffers::Inputs> explicitDraftTokensBuffers

std::optional<LookaheadDecodingBuffers> lookaheadOutputs

std::optional<EagleBuffers::Inputs> eagleBuffers

公共静态属性

static float constexpr kNegativeInfinity = -1e20f

class BeamHypotheses

公共函数

void empty(BufferManager const &manager)

void reshape(

SizeType32 batchSize,

SizeType32 beamWidth,

SizeType32 maxSequenceLength,

)

void release()

void init(BufferManager const &manager, TokenIdType endId)

BeamHypotheses slice(SizeType32 batchIndex, SizeType32 size) const

公共成员

TensorPtr outputIdsCBA

TensorPtr logProbsCBA

TensorPtr sequenceLengthsCBA

TensorPtr cumLogProbsCBA

TensorPtr normedScoresCBA

TensorPtr numBeamsCBA

TensorPtr minNormedScoresCBA

TensorPtr batchDones

class SpeculativeDecodingOutputs

公共成员

TensorPtr nextDraftTokens

TensorPtr nextDraftTokensLen

TensorPtr prevDraftTokensLen

TensorPtr acceptedTokensLen

TensorPtr acceptedLengthsCumSum

TensorPtr pathsOffsets

eagleBuffers.h

namespace tensorrt_llm

namespace batch_manager

namespace runtime

class EagleBuffers

公共类型

using LlmRequestPtr = std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>

using RequestVector = std::vector<LlmRequestPtr>

using SizeType32 = runtime::SizeType32

using ITensor = runtime::ITensor

using BufferPtr = runtime::IBuffer::SharedPtr

using TensorPtr = runtime::ITensor::SharedPtr

using TensorMap = runtime::StringPtrMap<runtime::ITensor>

公共函数

EagleBuffers(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

runtime::BufferManager const &manager,

runtime::ModelConfig const &modelConfig,

runtime::WorldConfig const &worldConfig,

executor::DecodingConfig const &decodingConfig,

)

void reshape(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

runtime::ModelConfig const &modelConfig,

)

void setFromInputs(

RequestVector const &contextRequests,

RequestVector const &genRequests,

runtime::ITensor const &requestTypes,

ITensor const &seqSlots,

EagleBuffers::Inputs const &decoderBuffers,

runtime::BufferManager const &manager,

runtime::ModelConfig const &modelConfig,

runtime::WorldConfig const &worldConfig,

) const

void insertInputTensors(

TensorMap &inputBuffers,

TensorMap &outputBuffers,

runtime::WorldConfig const &worldConfig,

) const

公共成员

Inputs engineInputs

class tensorrt_llm::runtime::EagleBuffers::EngineOutputs engineOutputs

私有函数

template<typename T>
void setFromInputs(

RequestVector const &contextRequests,

RequestVector const &genRequests,

SizeType32 vocabSizePadded,

ITensor const &seqSlots,

EagleBuffers::Inputs const &draftBuffers,

runtime::EagleModule const &eagleModule,

runtime::BufferManager const &manager,

) const

私有成员

std::size_t scanReduceTempStorageBytes = {0}

float mDefaultPosteriorThreshold = {0.09f}

bool mDoGreedySampling = {true}

BufferPtr scanReduceTempStorage

TensorPtr cumSumGenerationLengths

TensorPtr maxGenerationLength

TensorPtr chunkedContextNextTokensHost

TensorPtr greedySamplingHost

TensorPtr posteriorAlphaHost

TensorPtr posteriorThresholdHost

class EngineOutputs

公共成员

TensorPtr nextDraftTokens

[batchSize, 最大草稿解码词元数]

TensorPtr nextDraftLens

[批处理大小]

TensorPtr nextDraftPaths

[batchSize, 最大路径数, 最大路径长度]

TensorPtr acceptedTokens

[batchSize, 最大路径长度]

TensorPtr acceptedLens

[批处理大小]

TensorPtr acceptedPaths

[批处理大小]

TensorPtr chunkedContextNextTokens

[批处理大小]

class Inputs

公共函数

void create(

SizeType32 maxNumSequences,

BufferManager const &manager,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

)

公共成员

TensorPtr temperatures

[最大批处理大小] 或 [序列数]

TensorPtr posteriorAlpha

[最大批处理大小] 或 [序列数]

TensorPtr posteriorThreshold

[最大批处理大小] 或 [序列数]

TensorPtr randomDataSample

[最大批处理大小] 或 [序列数]

TensorPtr randomDataValidation

[最大批处理大小, 最大解码词元数] 或 [序列数, 最大解码词元数]

TensorPtr draftTokens

[最大批处理大小, 最大草稿解码词元数] 或 [序列数, 最大草稿解码词元数]

TensorPtr draftLens

[最大批处理大小] 或 [序列数]

TensorPtr draftPaths

[最大批处理大小, 最大路径数, 最大路径长度] 或 [序列数, 最大路径数, 最大路径长度]

TensorPtr specDecodingGenerationLengths

[最大批量大小] 或 [生成序列数量]

TensorPtr specDecodingGenerationLengthsHost

[最大批量大小] 或 [生成序列数量]

TensorPtr specDecodingPackedMasks

[最大批量大小, 最大解码令牌数, 向上取整(最大解码令牌数 / 32)] 或 [生成序列数量, 最大解码令牌数, 向上取整(最大解码令牌数 / 32)]

TensorPtr specDecodingPositionOffsets

[最大批量大小] 或 [生成序列数量]

TensorPtr eagleNetCtxRequestTypesHost

[最大批处理大小] 或 [序列数]

TensorPtr eagleNetCtxContextLengthsHost

[最大批处理大小] 或 [序列数]

TensorPtr eagleNetCtxPastKeyValueLengthsHost

[最大批处理大小] 或 [序列数]

TensorPtr eagleNetGenRequestTypesHost

[最大批处理大小] 或 [序列数]

TensorPtr eagleNetGenContextLengthsHost

[最大批处理大小] 或 [序列数]

TensorPtr eagleNetGenPastKeyValueLengthsHost

[最大批处理大小] 或 [序列数]

TensorPtr inputGenTokensHost

[最大批量大小 * 最大解码令牌数] 或 [序列数量 * 最大解码令牌数]

TensorPtr chunkedContextNextTokens

[最大批处理大小] 或 [序列数]

TensorPtr useSpecDecoding

[1]

TensorPtr useDynamicTreeHost

[1]

TensorPtr dynamicTreeMaxTopKHost

[1]

TensorPtr prevScores

[最大批处理大小, 最大草稿解码词元数] 或 [序列数, 最大草稿解码词元数]

TensorPtr currentExpandIndices

[最大批处理大小, 最大草稿解码词元数] 或 [序列数, 最大草稿解码词元数]

TensorPtr allLayersScores

[最大批量大小, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数] 或 [序列数量, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数]

TensorPtr allLayersDraftTokenIds

[最大批量大小, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数] 或 [序列数量, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数]

TensorPtr allLayersDraftTokenIdsPredecessor

[最大批量大小, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数] 或 [序列数量, Eagle层数, 最大解码草稿令牌数 * 最大解码草稿令牌数]

explicitDraftTokensBuffers.h

namespace tensorrt_llm

namespace runtime

class ExplicitDraftTokensBuffers

公共类型

using SizeType32 = runtime::SizeType32

using ITensor = runtime::ITensor

using BufferPtr = runtime::IBuffer::SharedPtr

using TensorPtr = runtime::ITensor::SharedPtr

using TensorMap = runtime::StringPtrMap<runtime::ITensor>

公共函数

ExplicitDraftTokensBuffers(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

runtime::BufferManager const &manager,

runtime::ModelConfig const &modelConfig,

runtime::WorldConfig const &worldConfig,

)

void reshape(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

runtime::ModelConfig const &modelConfig,

)

void setFromInputs(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

runtime::ITensor const &requestTypes,

ITensor const &seqSlots,

ExplicitDraftTokensBuffers::Inputs const &decoderBuffers,

ITensor const &contextPositionIds,

runtime::ModelConfig const &modelConfig,

runtime::WorldConfig const &worldConfig,

runtime::BufferManager const &manager,

runtime::CudaStream const &stream,

) const

void insertInputTensors(

TensorMap &inputBuffers,

TensorMap &outputBuffers,

runtime::WorldConfig const &worldConfig,

) const

公共成员

tensorrt_llm::runtime::ExplicitDraftTokensBuffers::EngineInputs engineInputs

class tensorrt_llm::runtime::ExplicitDraftTokensBuffers::EngineOutputs engineOutputs

std::size_t scanTempStorageBytes = {0}

BufferPtr scanTempStorage

TensorPtr cumSumGenerationLengths

私有函数

template<typename T>
void setFromInputs(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

SizeType32 vocabSizePadded,

ITensor const &seqSlots,

ExplicitDraftTokensBuffers::Inputs const &draftBuffers,

ITensor const &contextPositionIds,

runtime::ExplicitDraftTokensModule const &explicitDraftTokensModule,

runtime::CudaStream const &stream,

) const

class EngineInputs : public tensorrt_llm::runtime::ExplicitDraftTokensBuffers::Inputs

公共成员

TensorPtr requestTypesDevice

[序列数量]，在GPU上

TensorPtr positionOffsets

[生成序列数量]

class EngineOutputs

公共成员

TensorPtr nextGenerationLengths

[批处理大小]

TensorPtr nextPositionOffsets

[批处理大小]

TensorPtr masks

[批量大小, 最大解码令牌数, 最大解码令牌数]，布尔类型

TensorPtr nextDraftTokens

[batchSize, 最大路径数, 最大路径长度]

TensorPtr nextDraftIndices

[batchSize, 最大路径数, 最大路径长度]

TensorPtr nextDraftProbs

[批量大小, 最大路径数量, 最大草稿路径长度, 词汇表大小]

TensorPtr nextFlatTokens

[batchSize * maxDecodingTokens]

TensorPtr bestPathLengths

[批处理大小]

TensorPtr bestPathIndices

[批处理大小]

TensorPtr maxGenToken

[1]

TensorPtr totalGenToken

[1]

TensorPtr packedPositionIds

[batchSize * maxDecodingTokens]

class Inputs

由 tensorrt_llm::runtime::ExplicitDraftTokensBuffers::EngineInputs 继承

公共函数

void create(

SizeType32 maxNumSequences,

runtime::BufferManager const &manager,

runtime::ModelConfig const &modelConfig,

runtime::WorldConfig const &worldConfig,

)

公共成员

TensorPtr temperatures

[maxBatchSize]

TensorPtr positionIdsBase

[maxBatchSize]

TensorPtr generationLengths

[最大批量大小] 或 [生成序列数量]

TensorPtr randomDataSample

[maxBatchSize]

TensorPtr randomDataValidation

[maxBatchSize, maxNumPaths, maxPathDraftLen] 或 [numGenSequences, maxNumPaths, maxPathDraftLen]

TensorPtr draftTokens

[maxBatchSize, maxNumPaths, maxPathLen] 或 [numGenSequences, maxNumPaths, maxPathLen]

TensorPtr draftIndices

[maxBatchSize, maxNumPaths, maxPathLen] 或 [numGenSequences, maxNumPaths, maxPathLen]

TensorPtr draftProbs

[maxBatchSize, maxNumPaths, maxPathDraftLen, vocabSize] 或 [numGenSequences, maxNumPaths, maxPathDraftLen, vocabSize]

TensorPtr packedMasks

[最大批量大小, 最大解码令牌数, 向上取整(最大解码令牌数 / 32)] 或 [生成序列数量, 最大解码令牌数, 向上取整(最大解码令牌数 / 32)]

TensorPtr positionIds

[最大批量大小] 或 [生成序列数量]

TensorPtr maxGenLengthHost

TensorPtr generationLengthsHost

TensorPtr useSpecDecoding

generationInput.h

namespace tensorrt_llm

namespace runtime

class GenerationInput : public tensorrt_llm::runtime::GenericGenerationInput<ITensor::SharedPtr, PromptTuningParams>

公共类型

using Base = GenericGenerationInput<ITensor::SharedPtr, PromptTuningParams>

using TensorPtr = Base::TensorPtr

公共函数

inline explicit GenerationInput(

SizeType32 const endId,

SizeType32 const padId,

TensorPtr ids,

TensorPtr lengths,

bool packed = false,

)

template<typename TTensor, typename PromptTuningParams>
class GenericGenerationInput

#include <generationInput.h>

• endId，是标记输入序列结束的 token ID（也称为 EOS 或序列结束）。例如，对于词汇量为 50,257 个 token 的 GPT2 模型，它是 50,256，

• padId，是用于填充的 token ID（即，对于填充序列，填充输入长度或之后位置的槽位）。它可以设置为与 endId 相同的值，

• ids，是输入 ID 的张量。该张量必须分配在 GPU 上。当输入张量被填充时，ids 的形状是 [batchSize, maxInputLength]，其中 batchSize 和 maxInputLength 必须遵守传递给 GptSession 构造函数的 sessionConfig 中的最大尺寸。当输入被打包时，ids 的形状是 [numTokens]，其中 numTokens 是批处理中不同序列长度的总和，

• lengths，是输入序列长度的张量。该张量必须分配在 GPU 上并包含 batchSize 个值，

• packed，指示 ids 张量是打包的还是填充的。在此版本中，该标志必须与通过 ModelConfig 类的实例传递给构造函数的值匹配。在未来的版本中，session 可能会变得更加灵活，并自动填充或打包输入，

• embeddingBiasOpt，是 GPU 上的一个浮点值张量，包含在采样期间添加到 logits 的偏差（作为模型的最后一步，在从隐藏状态到 logits 的投影之后）。该张量必须包含 vocabSize 个元素（如传递给构造函数的 modelConfig 参数中所定义），

• badWordsList，是 GPU 上的一个整数张量，编码了必须从生成序列中禁止的单词列表。其形状如以下所述为 [2, badWordsLength]，或者当批处理中每个序列有不同列表时为 [batchSize, 2, badWordsLength]，

• stopWordsList，是 GPU 上的一个整数张量，编码了触发序列生成结束的单词列表。其形状如以下所述为 [2, stopWordsLength]，或者当批处理中每个序列有不同列表时为 [batchSize, 2, stopWordsLength]，

• maxNewTokens，是要生成的最大 token 数量。

badWordsList 和 stopWordsList 张量具有相同的形状 [2, length]。我们以一个包含三个单词的示例来说明这些列表的表示。第一个单词包含 token [5, 7, 3]，第二个包含 [9, 2]，第三个由 token [6, 2, 4, 1] 组成。总共有 9 个 token。这就是长度。张量的形状是 [2, 9]。张量的第一行必须包含这 9 个 token ID，第二行必须存储单词长度的 包容性前缀和，如下图所示

   0           3       5              9
   |           |       |              |
   V           V       V              V
[  5,  7,  3,  9,  2,  6,  2,  4,  1]
[  3,  5,  9, -1, -1, -1, -1, -1, -1]

如果所有单词都只包含一个 token，则张量的最内层维度必须增加 1（例如，对于由单个 token 组成的 4 个单词，长度必须是 5 而不是 4 — 形状是 [2, 5]）。

公共类型

using TensorPtr = TTensor

公共函数

inline explicit GenericGenerationInput(

SizeType32 const endId,

SizeType32 const padId,

TensorPtr ids,

TensorPtr lengths,

bool packed = false,

)

公共成员

SizeType32 endId

强制参数。

SizeType32 padId

TensorPtr ids

TensorPtr lengths

bool packed

TensorPtr embeddingBias

可选参数。

TensorPtr badWordsList

TensorPtr stopWordsList

std::optional<SizeType32> maxNewTokens

PromptTuningParams promptTuningParams

Ptuning 参数。

generationOutput.h

namespace tensorrt_llm

namespace runtime

class GenerationOutput : public tensorrt_llm::runtime::GenericGenerationOutput<ITensor::SharedPtr>

公共类型

using Base = GenericGenerationOutput<ITensor::SharedPtr>

using TensorPtr = Base::TensorPtr

公共函数

inline explicit GenerationOutput(TensorPtr ids, TensorPtr lengths)

template<typename TTensor>
class GenericGenerationOutput

#include <generationOutput.h>

• ids，是包含输出 token ID 的张量。其形状为 [batchSize, beamWidth, maxSeqLength]，其中 maxSeqLength 是 maxInputLength 和 maxNewTokens 之和。生成后，它包含每个序列的输入 token 的副本，后跟输出 token。当序列短于 maxSeqLength 时，会在序列末尾添加填充 token。

注意，此版本的 TensorRT-LLM 中该张量的形状与其先前版本中的形状不同，先前版本中为.

• logProbs，是 GPU 上用于存储生成 token 的对数概率 (log-prob) 的浮点值张量。其形状为 [maxNewTokens, batchSize, beamWidth]。其形状在未来版本中可能会更改，以匹配输出 ids 张量的形状。

• contextLogits，是 GPU 上用于存储 context logits 的值张量（数据类型与计算类型相同）。其形状为 [batchSize, maxSequenceLength, vocabSizePadded]。如果使用 remove_input_padding，其形状为 [packedSize, vocabSizePadded]。仅当 TensorRT 引擎构建时启用了 gather_context_logits 或 gather_all_token_logits 参数，此缓冲区才会被填充。

  推理完成后，您可以在 GenerationOutput.contextLogits 中获取 context logits，这些是 GPU 上的变量。具体获取方法请参考 gptSessionBenchmark.cpp 的示例。

  需要指出的是，启用此计算可能会影响性能（语言建模头部 (LM head) 必须对所有 context token 执行矩阵乘法，而不仅仅是对最后一个 token）。

• generationLogits，是 GPU 上用于存储生成 logits 的值张量（数据类型与计算类型相同）。其形状为 [batchSize, beamWidth, maxOutputLen, vocabSizePadded]。仅当 TensorRT 引擎构建时启用了 gather_generation_logits 或 gather_all_token_logits 参数，此缓冲区才会被填充。

  生成 logits 也可以在推理完成后通过 GenerationOutput.generationLogits 获取。

• onTokenGenerated，是在生成循环中调用的回调函数，用于在循环继续执行的同时将新生成的 token 传递给调用者。该回调函数的实现必须接受输出 ids 张量、生成 step 以及一个布尔标志，该标志指示生成是否完成。

公共类型

using TensorPtr = TTensor

using Callback = std::function<void(TensorPtr const &ids, SizeType32 step, bool finished)>

公共函数

inline explicit GenericGenerationOutput(

TensorPtr ids,

TensorPtr lengths,

)

公共成员

TensorPtr ids

强制参数。

TensorPtr lengths

TensorPtr cumLogProbs

可选参数。

TensorPtr logProbs

TensorPtr contextLogits

TensorPtr generationLogits

Callback onTokenGenerated

回调函数。

gptDecoderBatched.h

namespace tensorrt_llm

namespace batch_manager

namespace runtime

class GptDecoderBatched : public tensorrt_llm::runtime::IGptDecoderBatched

#include <gptDecoderBatched.h>

支持动态批处理的 GPT 解码器类。

公共类型

using CudaStreamPtr = std::shared_ptr<CudaStream>

using LlmRequestPtr = std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>

using RequestVector = std::vector<LlmRequestPtr>

using TensorPtr = ITensor::SharedPtr

using SharedConstPtr = ITensor::SharedConstPtr

公共函数

GptDecoderBatched(

CudaStreamPtr stream,

SpeculativeDecodingMode const &speculativeDecodingMode,

nvinfer1::DataType dtype,

)

virtual void setup(

executor::DecodingMode const &mode,

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

SizeType32 maxTokensPerStep,

nvinfer1::DataType dtype,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

) override

在调用 forward() 之前设置解码器

virtual void disableLookahead(

RequestVector const &genRequests,

TensorPtr const &batchSlots,

) override

禁用 Lookahead 解码。

virtual CudaEvent forwardAsync(

decoder_batch::Output &output,

decoder_batch::Input const &input,

) override

对所有请求运行一个步骤，不阻塞主机进程，并返回事件以进行同步。

virtual void forward(

decoder_batch::Output &output,

decoder_batch::Input const &input,

) override

对所有请求运行一个步骤，并在主机上等待完成。

virtual CudaEvent finalize(

decoder::DecoderState const &decoderState,

SizeType32 batchSlot,

SamplingConfig const &samplingConfig,

bool streaming,

) const override

收集请求 batchSlot 的最终 Beam Search 结果。结果仅在返回事件后可用。

inline decoder::DecoderState &getDecoderState()

inline decoder::DecoderState const &getDecoderState() const

inline CudaStreamPtr getDecoderStream() const

inline IGptDecoder &getUnderlyingDecoder() const

inline BufferManager const &getBufferManager() const

私有类型

using GptDecoderPtr = std::unique_ptr<IGptDecoder>

私有函数

void setExplicitDraftTokensInputs(decoder_batch::Input const &input)

设置显式草稿 token 的输入。

void setEagleInputs(decoder_batch::Input const &input)

设置用于 eagle 解码的输入。

void forwardDispatch(

decoder_batch::Output &output,

decoder_batch::Input const &input,

)

为每个引擎步骤调用解码器以处理令牌。

void prepareForward(

SizeType32 step,

decoder_batch::Output &output,

decoder_batch::Input const &input,

)

准备解码器步骤的输入和输出。

私有成员

CudaStreamPtr mRuntimeStream

CudaStreamPtr mDecoderStream

BufferManager mBufferManager

GptDecoderPtr mDecoder

std::shared_ptr<decoder::DecoderState> mDecoderState

gptSession.h

namespace tensorrt_llm

namespace batch_manager

namespace kv_cache_manager

namespace runtime

class GptSession

公共类型

using LoggerPtr = std::shared_ptr<nvinfer1::ILogger>

公共函数

GptSession(

Config const &sessionConfig,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

RawEngine const &rawEngine,

LoggerPtr logger = nullptr,

)

参数:

• sessionConfig – 会话配置，

• modelConfig – 模型描述，

• worldConfig – 环境描述，

• rawEngine – 编译后的 TensorRT 引擎，

• logger – 可选的日志记录器。

inline GptSession(

Config const &sessionConfig,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

void const *engineBuffer,

std::size_t engineSize,

LoggerPtr logger = nullptr,

)

inline GptSession(

Config const &sessionConfig,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

std::vector<uint8_t> const &engineBuffer,

LoggerPtr logger = nullptr,

)

GptSession(

Config const &sessionConfig,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

std::string const &engineFile,

LoggerPtr logger = nullptr,

)

nvinfer1::ILogger &getLogger() const

BufferManager const &getBufferManager() const

BufferManager::CudaStreamPtr getRuntimeStreamPtr() const

inline ModelConfig const &getModelConfig() const

inline WorldConfig const &getWorldConfig() const

inline int getDevice() const noexcept

inline bool getNormalizeLogProbs() const noexcept

inline bool getGatherGenerationLogits() const

nvinfer1::IEngineInspector &getEngineInspector() const

nvinfer1::DataType getLogitDataType() const

nvinfer1::DataType getTensorDataType(std::string const &name) const

nvinfer1::Dims getTensorShape(std::string const &name) const

void generate(

GenerationOutput &outputs,

GenerationInput const &inputs,

SamplingConfig const &samplingConfig,

std::shared_ptr<GenerationProfiler> const generationProfiler = nullptr,

)

此函数执行生成循环。

给定要读取的输入张量和要填充的输出张量，该成员函数可以生成结果，或者每个序列已达到完成状态（因为生成过程将运行生成循环，直到达到最大令牌数，或者遇到“序列结束”令牌或“停止词”列表中的某个词）。该函数的伪代码如下所示（成员函数名称已更改以保持说明简单）

// Have all the sequences in the batch reached completion?
bool allFinished = false;

// Until all sequences are finished or the number of steps reaches the limit...
for (int step = 0; !allFinished && step < maxNewTokens; ++step) {

// Trigger the computation of the logits...
computeLogits(...);

// Run the sampling to produce a token (for each active sequence) from the logits.
allFinished = generateTokensFromLogits(...);

// Callback to stream the output tokens while the generation loop continues.
onTokenGenerated(...);
}

void setLayerProfiler()

设置 LayerProfiler 以收集每层的性能信息。

std::string getLayerProfileInfo() const

打印每层的性能分析信息。

私有类型

using BaseKVCacheManager = batch_manager::kv_cache_manager::BaseKVCacheManager

using KvCacheConfig = batch_manager::kv_cache_manager::KvCacheConfig

using TensorPtr = runtime::ITensor::SharedPtr

using TokenGeneratedCallback = std::function<void(SizeType32 step, bool finished)>

私有函数

inline bool useCudaGraphs()

void generateBatched(

std::vector<GenerationOutput> &microBatchesOutputs,

std::vector<GenerationInput> const &microBatchesInputs,

SamplingConfig const &samplingConfig,

TokenGeneratedCallback const &onTokenGenerated,

std::shared_ptr<GenerationProfiler> const generationProfiler,

)

void setup(Config const &sessionConfig)

void createContexts()

void createBuffers(SizeType32 numMicroBatches)

void createDecoders(

SizeType32 batchSize,

SizeType32 beamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

nvinfer1::DataType logitsType,

bool decoderPerRequest,

SizeType32 numMicroBatches,

executor::DecodingMode const &decodingMode,

)

void createKvCacheManager(

SizeType32 batchSize,

SizeType32 beamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

KvCacheConfig const &config,

)

void createCustomAllReduceWorkspace(

SizeType32 batchSize,

SizeType32 beamWidth,

SizeType32 maxSequenceLength,

)

void executeContextStep(

std::vector<GenerationInput> const &generationBatchesInputs,

std::vector<SizeType32> const &generationBatchesOffsets,

BaseKVCacheManager const *kvCacheManager,

)

SizeType32 executeGenerationStep(

SizeType32 step,

std::vector<GenerationInput> const &microBatchesInputs,

std::vector<GenerationOutput> &microBatchesOutputs,

std::vector<SizeType32> const &microBatchOffsets,

BaseKVCacheManager *kvCacheManager,

std::vector<bool> &microBatchesFinished,

)

void decoderStepAsync(

SizeType32 decoderStep,

SizeType32 microBatchId,

)

在最后一个 PP 级别执行解码器，在其他 PP 级别接收解码器输出。

bool shouldStopSync(

SizeType32 batchSize,

SizeType32 beamWidth,

SizeType32 microBatchId,

)

与解码器同步并返回 shouldStop 标志。

void finalize(

SizeType32 microBatchId,

SamplingConfig const &samplingConfig,

)

在最后一个 PP 级别收集最终输出 ID 和 log prob，并将其发送到第一个 PP 级别。

主机上的接收是异步的，因此在访问前需要进行同步。

void kvCacheAddSequences(

SizeType32 beamWidth,

SizeType32 microBatchId,

SizeType32 firstBatchIdx,

)

ITensor::SharedPtr initDecoder(

ITensor &outputIds,

GenerationInput const &inputs,

GenerationOutput const &outputs,

SamplingConfig const &samplingConfig,

SizeType32 microBatchId,

) const

填充 outputIds 并返回对 newTokens 张量的引用。

TokenGeneratedCallback createOnTokenGeneratedCallback(

GenerationOutput &outputs,

)

bool shouldUseKVCacheManager() const

私有成员

ModelConfig const mModelConfig

WorldConfig const mWorldConfig

int mDevice = {-1}

std::shared_ptr<NcclCommunicator> mPipelineComm

std::shared_ptr<CudaStream> mCommStream

CudaEvent mCommEvent = {}

std::shared_ptr<AllReduceBuffers> mAllReduceBuffers

SizeType32 mDecoderMaxSequenceLength = {}

std::vector<SizeType32> mDecoderMaxAttentionWindowVec = {}

SizeType32 mDecoderMaxAttentionWindow = {}

SizeType32 mDecoderSinkTokenLength = {}

LoggerPtr mLogger

std::shared_ptr<TllmRuntime> mRuntime

std::shared_ptr<BaseKVCacheManager> mKvCacheManager

MicroBatchConfig mMicroBatchConfig

std::vector<std::shared_ptr<IStatefulGptDecoder>> mDecoders

std::vector<std::shared_ptr<RuntimeBuffers>> mBuffers

std::vector<CudaEvent> mReceivedEvents

bool mCudaGraphMode = {false}

std::vector<CudaGraphExecutor> mCudaGraphInstances

bool mNormalizeLogProbs = {true}

bool mGatherGenerationLogits = {false}

友元

friend class batch_manager::TrtGptModelV1

class Config

#include <gptSession.h>

会话执行和缓冲区大小的配置。可以在调用 generate 时使用的批量大小和束搜索宽度小于配置的参数。

maxBatchSize 将被微批量数量划分，以初始化每个批量缓冲区。

公共函数

inline Config(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxSequenceLength,

float gpuWeightsPercent = 1.0,

)

公共成员

SizeType32 maxBatchSize

SizeType32 maxBeamWidth

SizeType32 maxSequenceLength

bool useGpuDirectStorage = {false}

float gpuWeightsPercent

bool decoderPerRequest = {false}

bool cudaGraphMode = {false}

KvCacheConfig kvCacheConfig = {}

std::optional<SizeType32> ctxMicroBatchSize = std::nullopt

std::optional<SizeType32> genMicroBatchSize = std::nullopt

std::optional<executor::DecodingMode> decodingMode = std::nullopt

bool normalizeLogProbs = true

bool gatherGenerationLogits = false

class CudaGraphExecutor

公共函数

CudaGraphExecutor() = default

inline ~CudaGraphExecutor()

inline bool hasInstance()

void clear()

void prepareNextGraph(

TllmRuntime const &runtime,

SizeType32 nextContextId,

)

void launch(CudaStream const &stream)

私有函数

void create(cudaGraph_t const &graph)

bool update(cudaGraph_t const &graph)

void uploadToStream(CudaStream const &stream)

私有成员

cudaGraphExec_t mInstance

class GenerationProfiler

#include <gptSession.h>

用于对推理请求的生成阶段进行性能分析的可选分析器类。

公共函数

inline GenerationProfiler()

inline CudaEvent const &getStart() const

inline CudaEvent const &getEnd() const

inline float getElapsedTimeMs()

公共静态属性

static constexpr unsigned int flags = {cudaEventDefault}

私有成员

CudaEvent start

CudaEvent end

class MicroBatchConfig

公共函数

inline MicroBatchConfig()

explicit MicroBatchConfig(

SizeType32 maxBatchSize,

SizeType32 pipelineParallelism,

std::optional<SizeType32> genMicroBatchSize,

std::optional<SizeType32> ctxMicroBatchSize,

)

inline constexpr SizeType32 numCtxPerGen() const

inline constexpr SizeType32 getGenGraphId(

SizeType32 flipFlopId,

SizeType32 generationBatchId,

) const

在每个生成批次之间切换使用两个图实例。

公共成员

SizeType32 numCtxBatches

SizeType32 numGenBatches

SizeType32 ctxBatchSize

SizeType32 genBatchSize

namespace utils

函数

std::vector<uint8_t> loadEngine(std::string const &enginePath)

gptDecoder.h

namespace tensorrt_llm

namespace layers

namespace runtime

函数

inline runtime::ITensor::SharedConstPtr getDefaultBatchSlots(

runtime::SizeType32 batchSize,

)

用于生成批次槽 [0, 1, ..., batchSize - 1] 的辅助函数，适用于未向解码器明确提供批次槽的路径。

template<typename T>
class GptDecoder : public virtual tensorrt_llm::runtime::IGptDecoder

公共类型

using CudaStreamPtr = BufferManager::CudaStreamPtr

using TensorPtr = std::shared_ptr<ITensor>

公共函数

GptDecoder(

executor::DecodingMode const &mode,

size_t maxBatchSize,

size_t maxBeamWidth,

size_t vocabSize,

size_t vocabSizePadded,

size_t maxSequenceLength,

CudaStreamPtr const &stream,

std::shared_ptr<SpeculativeDecodingModule const> speculativeDecodingModule = nullptr,

)

virtual void setup(

SamplingConfig const &samplingConfig,

size_t batchSize,

TensorConstPtr const &batchSlots,

std::optional<DecodingOutput> const &output = std::nullopt,

std::optional<std::vector<decoder_batch::Request> const> const &requests = std::nullopt,

) override

virtual void forwardAsync(

DecodingOutput &output,

DecodingInput const &input,

) override

virtual void forwardSync(

DecodingOutput &output,

DecodingInput const &input,

) override

inline virtual SamplingConfig const &getSamplingConfig() const override

virtual void disableLookahead(

std::optional<SamplingConfig> const &samplingConfig,

SizeType32 batchSize,

TensorConstPtr batchSlots,

) override

私有成员

std::shared_ptr<BufferManager> mManager

std::shared_ptr<tensorrt_llm::layers::DynamicDecodeLayer<T>> mDynamicDecodeLayer

std::shared_ptr<tensorrt_llm::runtime::DecodingLayerWorkspace> mDecodingLayerWorkspace

SamplingConfig mSamplingConfig

size_t mMaxBatchSize

size_t mVocabSize

size_t mVocabSizePadded

executor::DecodingMode mDecodingMode

class IGptDecoder

被子类 tensorrt_llm::runtime::GptDecoder< T > 继承

公共类型

using TensorPtr = runtime::ITensor::SharedPtr

using TensorConstPtr = runtime::ITensor::SharedConstPtr

公共函数

virtual ~IGptDecoder() = default

virtual void setup(

SamplingConfig const &samplingConfig,

size_t batchSize,

TensorConstPtr const &batchSlots,

std::optional<DecodingOutput> const &output = std::nullopt,

std::optional<std::vector<decoder_batch::Request> const> const &requests = std::nullopt,

) = 0

virtual void forwardAsync(

DecodingOutput &output,

DecodingInput const &input,

) = 0

virtual void forwardSync(

DecodingOutput &output,

DecodingInput const &input,

) = 0

virtual SamplingConfig const &getSamplingConfig() = 0

virtual void disableLookahead(

std::optional<SamplingConfig> const &samplingConfig,

SizeType32 batchSize,

TensorConstPtr batchSlots,

) = 0

公共静态函数

static inline std::unique_ptr<IGptDecoder> create(

executor::DecodingMode const &mode,

nvinfer1::DataType dtype,

size_t maxBatchSize,

size_t maxBeamWidth,

size_t vocabSize,

size_t vocabSizePadded,

size_t maxSequenceLength,

BufferManager::CudaStreamPtr const &stream,

std::shared_ptr<SpeculativeDecodingModule const> const &speculativeDecodingModule = nullptr,

)

iGptDecoderBatched.h

namespace tensorrt_llm

namespace batch_manager

namespace runtime

class IGptDecoderBatched

#include <iGptDecoderBatched.h>

支持动态批处理的 GPT 解码器类。

被子类 tensorrt_llm::runtime::GptDecoderBatched 继承

公共类型

using CudaStreamPtr = std::shared_ptr<CudaStream>

using LlmRequestPtr = std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>

using RequestVector = std::vector<LlmRequestPtr>

using TensorPtr = std::shared_ptr<ITensor>

公共函数

virtual void setup(

executor::DecodingMode const &mode,

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

SizeType32 maxTokensPerStep,

nvinfer1::DataType dtype,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

) = 0

在调用 forward() 之前设置解码器

virtual void disableLookahead(

RequestVector const &genRequests,

TensorPtr const &batchSlots,

) = 0

禁用 Lookahead 解码。

virtual CudaEvent forwardAsync(

decoder_batch::Output &output,

decoder_batch::Input const &input,

) = 0

对所有请求运行一个步骤，不阻塞主机进程，并返回事件以进行同步。

virtual void forward(

decoder_batch::Output &output,

decoder_batch::Input const &input,

) = 0

对所有请求运行一个步骤，并在主机上等待完成。

virtual CudaEvent finalize(

decoder::DecoderState const &decoderState,

SizeType32 batchSlot,

SamplingConfig const &samplingConfig,

bool streaming,

) const = 0

收集请求 batchIdx 的最终束搜索结果。结果只在事件返回后可用。

保护函数

IGptDecoderBatched() = default

virtual ~IGptDecoderBatched() = default

namespace decoder

namespace decoder_batch

class Input

公共类型

using TensorConstPtr = ITensor::SharedConstPtr

using TensorPtr = ITensor::SharedPtr

公共函数

inline explicit Input(

std::vector<std::vector<TensorConstPtr>> const &logits,

SizeType32 maxDecoderSteps,

)

inline explicit Input(std::vector<TensorConstPtr> const &logits)

公共成员

std::vector<std::vector<TensorConstPtr>> logits

强制参数。

[maxDecoderSteps][batchSize][1, beamWidth, vocabSizePadded]，在 GPU 上

SizeType32 maxDecoderSteps

活动 slot 的解码 token 的最大数量。

std::vector<TensorPtr> batchSlots

按槽排序的活跃解码器槽批次，[maxDecoderSteps][batchSize]。

TensorPtr batchSlotsRequestOrder

按请求顺序填充的槽，[batchSize]。

TensorPtr cacheIndirection

用于束搜索，同一个束内不同射线的KV缓存索引。

std::vector<std::vector<TensorPtr>> predictedDraftLogits

[maxBatchSize][maxAcceptedDraftTokensPerStep][maxDraftTokens + 1, vocabSizePadded]

std::optional<ExplicitDraftTokensBuffers::EngineOutputs> explicitDraftTokensInputs

显式草稿令牌数据。

std::optional<ExplicitDraftTokensBuffers::EngineInputs> explicitDraftTokensLastInputs

std::optional<EagleBuffers::EngineOutputs> eagleInputs

std::optional<EagleBuffers::EngineOutputs> eagleInputs

Eagle 数据。

std::optional<EagleBuffers::Inputs> eagleLastInputs

公共类型

using TensorPtr = std::shared_ptr<ITensor>

公共函数

class Output

公共成员

using TensorPtr = std::shared_ptr<ITensor>

Output() = default

namespace tensorrt_llm

namespace runtime

TensorPtr cacheIndirection

公共函数

gptJsonConfig.h

class GptJsonConfig,

inline GptJsonConfig(,

std::string name,

std::string version,

SizeType32 pipelineParallelism,

std::string precision,

SizeType32 tensorParallelism,

SizeType32 contextParallelism,

SizeType32 gpusPerNode,

)

ModelConfig modelConfig

std::optional<RuntimeDefaults> runtimeDefaults = std::nullopt

inline ModelConfig const &getModelConfig() const

inline ModelConfig &getModelConfigMutable()

inline std::string const &getName() const

inline std::string const &getVersion() const

inline std::string const &getPrecision() const

inline SizeType32 constexpr getTensorParallelism() const

inline SizeType32 constexpr getPipelineParallelism() const

inline SizeType32 constexpr getContextParallelism() const

inline SizeType32 constexpr getGpusPerNode() const

inline SizeType32 constexpr getWorldSize() const

WorldConfig const &worldConfig,

inline std::optional<RuntimeDefaults> getRuntimeDefaults() const,

std::string engineFilename(

std::string const &model

WorldConfig const &worldConfig,

) const

公共静态函数

inline std::string engineFilename(

) const

static GptJsonConfig parse(std::string const &json)

私有成员

static GptJsonConfig parse(std::istream &json)

static GptJsonConfig parse(std::filesystem::path const &path)

std::string const mName

std::string const mVersion

std::string const mPrecision

SizeType32 const mTensorParallelism

SizeType32 const mPipelineParallelism

SizeType32 const mContextParallelism

SizeType32 const mGpusPerNode

ModelConfig mModelConfig

namespace tensorrt_llm

namespace batch_manager

namespace runtime

std::optional<RuntimeDefaults> mRuntimeDefaults

iStatefulGptDecoder.h

支持动态批处理的 GPT 解码器类。

class IStatefulGptDecoder

公共类型

#include <iStatefulGptDecoder.h>

继承自 tensorrt_llm::runtime::StatefulGptDecoderBatched

公共函数

using CudaStreamPtr = std::shared_ptr<CudaStream>

executor::DecodingMode const &mode,

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

nvinfer1::DataType dtype,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

) = 0

using TensorPtr = std::shared_ptr<ITensor>

virtual void setup(

GenerationInput const &inputs,

GenerationOutput const &outputs,

SamplingConfig const &samplingConfig,

ModelConfig const &modelConfig,

) = 0

在调用 forward() 之前设置解码器，也调用 reshapeBuffers。

virtual void forwardAsync(

decoder::Output &output,

decoder::Input const &input,

) = 0

在不阻塞主机线程的情况下，为所有请求运行一步。

virtual void forwardSync() = 0

等待对 forwardAsync 的最后一次调用完成。

inline virtual void forward(

decoder::Output &output,

decoder::Input const &input,

)

为所有请求运行一步。

virtual void finalize(SamplingConfig const &samplingConfig) const = 0

收集所有请求的最终束搜索结果。

virtual TensorPtr getIds() const = 0

返回：

[batchSize, beamWidth, maxSequenceLength]，所有 token id，在 GPU 上

virtual TensorPtr getGatheredIds() const = 0

返回：

[batchSize, beamWidth, maxSequenceLength] 在 gatherTree 后的 token id

virtual TensorPtr getCumLogProbs() const = 0

返回：

[batchSize, maxBeamWidth]，累积对数概率（每个 beam），在 GPU 上

virtual TensorPtr getLogProbs() const = 0

返回：

[batchSize, maxBeamWidth, maxSequenceLength]，对数概率（每个 beam），在 GPU 上

virtual TensorPtr getNewTokens(SizeType32 iter = 0) const = 0

获取上一次 forward 传递中一步生成的 token。

参数:

**iter** – 要获取 token 的迭代次数，范围在 [0; maxTokensPerStep) 内

返回：

[batchSize, beamWidth]，在 iter 中生成的 token（每个 beam），在 GPU 上

virtual TensorPtr getNbFinished() const = 0

返回：

[1]，已完成序列的数量，在 pinned 主机内存中

virtual ~IStatefulGptDecoder() = default

保护函数

IStatefulGptDecoder() = default

namespace decoder

class Input

公共类型

using TensorPtr = ITensor::SharedPtr

公共函数

inline explicit Input(TensorPtr logits)

公共成员

TensorPtr logits

TensorPtr cacheIndirection

class Output

公共类型

using TensorPtr = std::shared_ptr<ITensor>

公共函数

Output() = default

公共成员

TensorPtr cacheIndirection

TensorPtr sequenceLengths

iBuffer.h

namespace tensorrt_llm

namespace runtime

类型定义

template<typename T>
using PointerElementType = typename std::remove_reference_t<T>::element_type

枚举

enum class MemoryType : std::int32_t

值

enumerator kGPU

enumerator kCPU

enumerator kPINNED

enumerator kUVM

enumerator kPINNEDPOOL

函数

template<typename T>
std::shared_ptr<std::remove_const_t<T>> constPointerCast(

std::shared_ptr<T> const &ptr,

) noexcept

template<typename T, typename D>
std::shared_ptr<std::remove_const_t<T>> constPointerCast(

std::unique_ptr<T, D> &&ptr,

) noexcept

template<typename T>
T const *bufferCast(IBuffer const &buffer)

获取指向缓冲区底层常量数据的类型化指针。

模板参数：

**T** – 底层数据的类型。

参数:

**buffer** – 要获取指针的缓冲区。

返回：

指向常量 T 的指针。

template<typename T>
T *bufferCast(IBuffer &buffer)

获取指向缓冲区底层数据的类型化指针。

模板参数：

**T** – 底层数据的类型。

参数:

**buffer** – 要获取指针的缓冲区。

返回：

指向 T 的指针。

template<typename T>
T *bufferCastOrNull(

IBuffer::SharedPtr const &bufferPtr,

)

检索指向 bufferPtr 指向的缓冲区底层数据的 T 类型指针，如果 bufferPtr 为空则返回 nullptr。

模板参数：

**T** – 底层数据的类型。

参数:

**bufferPtr** – 可能为空的共享指针。

返回：

指向 T 的指针，可能为 nullptr。

template<typename T>
T const *bufferCastOrNull(

IBuffer::SharedConstPtr const &bufferPtr,

)

检索指向 bufferPtr 指向的缓冲区底层常量数据的 T 类型指针，如果 bufferPtr 为空则返回 nullptr。

模板参数：

**T** – 底层数据的类型。

参数:

**bufferPtr** – 可能为空的共享指针。

返回：

指向常量 T 的指针，可能为 nullptr。

template<typename T>
T *bufferCastOrNull(

std::optional<IBuffer::SharedPtr> const &optionalBufferPtr,

)

检索指向 optionalBufferPtr 中包含的缓冲区指针指向的缓冲区底层数据的 T 类型指针，如果 optional 没有值，则返回 nullptr。

模板参数：

**T** – 底层数据的类型。

参数:

optionalBufferPtr – 可能为空的 optional 对象。

返回：

指向 T 的指针，可能为 nullptr。

template<typename T>
T const *bufferCastOrNull(

std::optional<IBuffer::SharedConstPtr> const &optionalBufferPtr,

)

检索指向 optionalBufferPtr 中包含的缓冲区指针指向的缓冲区底层数据的 T const 类型指针，如果 optional 没有值，则返回 nullptr。

模板参数：

**T** – 底层数据的类型。

参数:

optionalBufferPtr – 可能为空的 optional 对象。

返回：

指向常量 T 的指针，可能为 nullptr。

std::ostream &operator<<(

std::ostream &output,

IBuffer const &buffer,

)

用于打印缓冲区的实用函数。

class BufferDataType

#include <iBuffer.h>

围绕 nvinfer1::DataType 的包装器，提供对指针类型的支持。

公共函数

inline constexpr BufferDataType(

nvinfer1::DataType dataType,

bool _unsigned = false,

bool pointer = false,

)

inline constexpr operator nvinfer1::DataType() const noexcept

inline constexpr nvinfer1::DataType getDataType() const noexcept

inline constexpr bool isPointer() const noexcept

inline constexpr bool isUnsigned() const

inline constexpr std::size_t getSize() const noexcept

inline constexpr std::size_t getSizeInBits() const noexcept

公共静态属性

static auto constexpr kTrtPointerType = nvinfer1::DataType::kINT64

私有成员

nvinfer1::DataType mDataType

bool mUnsigned

bool mPointer

template<typename T>
class BufferRange : public tensorrt_llm::common::ArrayView<T>

公共类型

using Base = tensorrt_llm::common::ArrayView<T>

公共函数

inline BufferRange(T *data, size_type size)

template<typename U = T, std::enable_if_t<!std::is_const_v<U>, bool> = true>
inline explicit BufferRange(

IBuffer &buffer,

)

template<typename U = T, std::enable_if_t<std::is_const_v<U>, bool> = true>
inline explicit BufferRange(

IBuffer const &buffer,

)

template<nvinfer1::DataType kDataType, bool kIsUnsigned = false, bool kIsPointer = false>
struct DataTypeTraits

#include <iBuffer.h>

用于将 TensorRT 数据类型转换为 C++ 数据类型。

template<nvinfer1::DataType kDataType, bool kUnsigned>
struct DataTypeTraits<kDataType, kUnsigned, true>

公共类型

using type = typename DataTypeTraits<kDataType, kUnsigned, false>::type*

公共静态属性

static char constexpr name[] = "*"

static auto constexpr size = sizeof(type)

template<bool kUnsigned>
struct DataTypeTraits<nvinfer1::DataType::kBOOL, kUnsigned>

公共类型

using type = bool

公共静态属性

static char constexpr name[] = "bool"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kFLOAT>

公共类型

using type = float

公共静态属性

static char constexpr name[] = "float"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kHALF>

公共类型

using type = half

公共静态属性

static char constexpr name[] = "half"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kINT32>

公共类型

using type = std::int32_t

公共静态属性

static char constexpr name[] = "int32"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kINT32, true>

公共类型

using type = std::uint32_t

公共静态属性

static char constexpr name[] = "uint32"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kINT64>

公共类型

using type = std::int64_t

公共静态属性

static char constexpr name[] = "int64"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kINT64, true>

公共类型

using type = std::uint64_t

公共静态属性

static char constexpr name[] = "uint64"

static auto constexpr size = sizeof(type)

template<>
struct DataTypeTraits<nvinfer1::DataType::kINT8>

公共类型

using type = std::int8_t

公共静态属性

static char constexpr name[] = "int8"

static auto constexpr size = sizeof(type)

template<bool kUnsigned>
struct DataTypeTraits<nvinfer1::DataType::kUINT8, kUnsigned>

公共类型

using type = std::uint8_t

公共静态属性

static char constexpr name[] = "uint8"

static auto constexpr size = sizeof(type)

class IBuffer

继承自 tensorrt_llm::runtime::ITensor

公共类型

using UniquePtr = std::unique_ptr<IBuffer>

using SharedPtr = std::shared_ptr<IBuffer>

using UniqueConstPtr = std::unique_ptr<IBuffer const>

using SharedConstPtr = std::shared_ptr<IBuffer const>

using DataType = nvinfer1::DataType

公共函数

virtual void *data() = 0

返回指向底层数组的指针。

virtual void const *data() const = 0

返回指向底层数组的指针。

inline virtual void *data(std::size_t index)

返回指向底层数组中给定元素索引处的指针。

inline virtual void const *data(std::size_t index) const

返回指向底层数组中给定元素索引处的指针。

virtual std::size_t getSize() const = 0

返回缓冲区的大小（元素个数）。

inline virtual std::size_t getSizeInBytes() const

返回缓冲区的大小（以字节为单位）。

virtual std::size_t getCapacity() const = 0

返回缓冲区的容量。

virtual DataType getDataType() const = 0

返回缓冲区的数据类型。

virtual char const *getDataTypeName() const

virtual MemoryType getMemoryType() const = 0

返回缓冲区的内存类型。

virtual char const *getMemoryTypeName() const

virtual void resize(std::size_t newSize) = 0

调整缓冲区大小。如果新大小小于或等于当前容量，则此操作无效。

virtual void release() = 0

释放缓冲区。它将被重置为 nullptr。

virtual ~IBuffer() = default

IBuffer(IBuffer const&) = delete

不允许复制。

IBuffer &operator=(IBuffer const&) = delete

不允许复制。

公共静态函数

static char const *getDataTypeName(DataType dataType)

static UniquePtr slice(

SharedPtr buffer,

std::size_t offset,

std::size_t size,

)

在底层 buffer 上创建一个切片视图。该视图将具有与 buffer 相同的数据类型。

参数:

• buffer – 要查看的缓冲区。

• offset – 视图的偏移量。

• size – 视图的大小。

返回：

buffer 的一个视图。

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::size_t offset,

std::size_t size,

)

static inline UniquePtr slice(SharedPtr buffer, std::size_t offset)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::size_t offset,

)

static inline UniquePtr view(SharedPtr tensor)

返回底层 tensor 的一个视图，该视图可以独立调整大小。

参数:

tensor – 要查看的张量。

返回：

tensor 的一个视图。

static inline UniquePtr view(SharedPtr tensor, std::size_t size)

返回底层 tensor 的一个具有不同大小的视图。

参数:

• tensor – 要查看的张量。

• size – 视图的大小。

返回：

tensor 的一个视图。

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr view(

TConstPtr &&tensor,

std::size_t size,

)

static UniquePtr wrap(

void *data,

DataType type,

std::size_t size,

std::size_t capacity,

)

将给定的 data 封装在 IBuffer 中。IBuffer 不拥有底层 data，且不能调整大小超过 capacity。

参数:

• data – 要封装的数据。

• type – data 的数据类型。

• size – 缓冲区的大小。

• capacity – 缓冲区的容量。

返回：

一个 IBuffer。

static inline UniquePtr wrap(

void *data,

DataType type,

std::size_t size,

)

template<typename T>
static inline UniquePtr wrap(

T *data,

std::size_t size,

std::size_t capacity,

)

template<typename T>
static inline UniquePtr wrap(

T *data,

std::size_t size,

)

template<typename T>
static inline UniquePtr wrap(

std::vector<T> &v,

)

static MemoryType memoryType(void const *data)

确定指针的内存类型。

保护函数

IBuffer() = default

inline std::size_t toBytes(std::size_t size) const

返回以字节为单位的数组索引或大小。

template<MemoryType T>
struct MemoryTypeString

template<>
struct MemoryTypeString<MemoryType::kCPU>

公共静态属性

static auto constexpr value = "CPU"

template<>
struct MemoryTypeString<MemoryType::kGPU>

公共静态属性

static auto constexpr value = "GPU"

template<>
struct MemoryTypeString<MemoryType::kPINNED>

公共静态属性

static auto constexpr value = "PINNED"

template<>
struct MemoryTypeString<MemoryType::kPINNEDPOOL>

公共静态属性

static auto constexpr value = "PINNEDPOOL"

template<>
struct MemoryTypeString<MemoryType::kUVM>

公共静态属性

static constexpr value = "UVM"

template<typename T, bool = false>
struct TRTDataType

#include <iBuffer.h>

用于将 C++ 数据类型转换为 TensorRT 数据类型。

template<>
struct TRTDataType<bool>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kBOOL

template<>
struct TRTDataType<float>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kFLOAT

template<>
struct TRTDataType<half>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kHALF

template<>
struct TRTDataType<kernels::FinishedState>

公共静态属性

static constexpr auto value = TRTDataType<kernels::FinishedState::UnderlyingType>::value

template<>
struct TRTDataType<kernels::KVCacheIndex>

公共静态属性

static constexpr auto value = TRTDataType<kernels::KVCacheIndex::UnderlyingType>::value

template<>
struct TRTDataType<runtime::RequestType>

公共静态属性

static constexpr auto value = TRTDataType<std::underlying_type_t<runtime::RequestType>>::value

template<>
struct TRTDataType<std::int32_t>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kINT32

template<>
struct TRTDataType<std::int64_t>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kINT64

template<>
struct TRTDataType<std::int8_t>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kINT8

template<>
struct TRTDataType<std::uint32_t>

公共静态属性

static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT32, true}

template<>
struct TRTDataType<std::uint64_t>

公共静态属性

static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT64, true}

template<>
struct TRTDataType<std::uint8_t>

公共静态属性

static constexpr auto value = nvinfer1::DataType::kUINT8

template<typename T>
struct TRTDataType<T*>

公共静态属性

static auto constexpr value = BufferDataType{kUnderlyingType.getDataType(), kUnderlyingType.isUnsigned(), true}

私有静态属性

static auto constexpr kUnderlyingType = BufferDataType{TRTDataType<std::remove_const_t<T>, false>::value}

template<>
struct TRTDataType<void*>

公共静态属性

static constexpr auto value = BufferDataType::kTrtPointerType

ipcUtils.h

namespace tensorrt_llm

namespace runtime

函数

void lamportInitializeAll(

void *buffer_0,

void *buffer_1,

void *buffer_2,

size_t size,

)

bool canAccessPeer(WorldConfig const &worldConfig)

class AllReduceBuffers

公共类型

using TensorPtr = ITensor::SharedPtr

公共函数

AllReduceBuffers(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxSequenceLength,

SizeType32 hiddenSize,

BufferManager const &manager,

WorldConfig const &worldConfig,

bool const fakeBuffers = false,

)

公共成员

TensorPtr mAllReduceCommPtrs

std::vector<runtime::IpcMemory> mIpcMemoryHandles

class IpcMemory

公共类型

using BufferPtr = IBuffer::SharedPtr

公共函数

IpcMemory(

std::size_t bufferSize,

BufferManager const &manager,

WorldConfig const &worldConfig,

bool openIpc = true,

)

~IpcMemory()

IpcMemory(IpcMemory const&) = delete

IpcMemory &operator=(IpcMemory const&) = delete

IpcMemory(IpcMemory&&) = default

IpcMemory &operator=(IpcMemory&&) = default

inline std::vector<void*> const &getCommPtrs() const

公共静态属性

static size_t constexpr FLAGS_SIZE = (tensorrt_llm::kernels::MAX_ALL_REDUCE_BLOCKS + 1) * sizeof(uint32_t)

私有函数

void allocateIpcMemory(

std::size_t bufferSize,

BufferManager const &manager,

WorldConfig const &worldConfig,

)

void destroyIpcMemory()

私有成员

SizeType32 mTpRank

std::vector<void*> mCommPtrs

BufferPtr mBuffer

bool mOpenIpc

lookaheadBuffers.h

namespace tensorrt_llm

namespace runtime

class LookaheadDecodingBuffers

公共类型

using TensorPtr = ITensor::SharedPtr

公共函数

LookaheadDecodingBuffers(

SizeType32 maxNumSequences,

SizeType32 maxTokensPerStep,

BufferManager const &bufferManager,

)

公共成员

TensorPtr generationLengths

TensorPtr positionOffsets

TensorPtr packedMasks

TensorPtr positionIds

class LookaheadRuntimeBuffers

公共类型

using TensorPtr = ITensor::SharedPtr

using TensorMap = StringPtrMap<ITensor>

公共函数

LookaheadRuntimeBuffers(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

BufferManager const &manager,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

executor::DecodingConfig const &decodingConfig,

TllmRuntime const &runtime,

)

void setFromInputs(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

ITensor const &requestTypes,

ITensor const &seqSlots,

LookaheadDecodingBuffers const &decoderLookaheadBuffers,

TllmRuntime const &runtime,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

) const

void reshape(

SizeType32 numCtxSequences,

SizeType32 numGenSequences,

SizeType32 tokensPerStep,

)

void insertInputTensors(

TensorMap &inputBuffers,

TensorMap &outputBuffers,

WorldConfig const &worldConfig,

) const

void enableLookaheadDecoding(

SizeType32 maxBatchSize,

SizeType32 tokensPerStep,

)

void disableLookaheadDecoding()

公共成员

TensorPtr cumSumLength

TensorPtr packedMasksDevice

TensorPtr generationLengthsDevice

TensorPtr positionOffsetsDevice

TensorPtr positionIdsDevice

TensorPtr packedMaskHost

TensorPtr generationLengthsHost

TensorPtr positionOffsetsHost

TensorPtr positionIdsHost

TensorPtr packedMaskHostCopy

TensorPtr generationLengthsHostCopy

TensorPtr positionOffsetsHostCopy

TensorPtr positionIdsHostCopy

TensorPtr useSpecDecoding

TensorPtr batchSlotsHostCopy

iTensor.h

namespace nvinfer1

namespace tensorrt_llm

namespace runtime

函数

inline std::ostream &operator<<(

std::ostream &output,

ITensor::Shape const &dims,

)

用于打印形状的工具函数。

std::ostream &operator<<(

std::ostream &output,

ITensor const &tensor,

)

用于打印带形状的张量的工具函数。

template<typename T>
T const *bufferCastOrNull(

ITensor::SharedConstPtr const &tensorPtr,

)

检索指向 tensorPtr 所指向的张量底层数据的 T const 类型指针，如果 tensorPtr 为空则返回 nullptr。

必须声明此重载以避免涉及隐式转换为 IBuffer 时出现的歧义。

模板参数：

**T** – 底层数据的类型。

参数:

tensorPtr – 一个可能是 null 的共享指针。

返回：

指向 T const 的指针，可能为 nullptr。

template<typename T>
T *bufferCastOrNull(

ITensor::SharedPtr const &tensorPtr,

)

检索指向 tensorPtr 所指向的缓冲区底层数据的 T 类型指针，如果 tensorPtr 为空则返回 nullptr。

必须声明此重载以避免涉及隐式转换为 IBuffer 时出现的歧义。

模板参数：

**T** – 底层数据的类型。

参数:

tensorPtr – 一个可能是 null 的共享指针。

返回：

指向 T 的指针，可能为 nullptr。

template<typename T>
T *bufferCastOrNull(

std::optional<ITensor::SharedPtr> const &optionalTensorPtr,

)

检索指向 optionalBufferPtr 中包含的张量指针所指向的张量底层数据的 T 类型指针，如果 optional 没有值则返回 nullptr。

必须声明此重载以避免涉及隐式转换为 IBuffer 时出现的歧义。

模板参数：

**T** – 底层数据的类型。

参数:

optionalBufferPtr – 可能为空的 optional 对象。

返回：

指向 T 的指针，可能为 nullptr。

template<typename T>
T const *bufferCastOrNull(

std::optional<ITensor::SharedConstPtr> const &optionalTensorPtr,

)

检索指向 optionalBufferPtr 中包含的张量指针所指向的张量底层数据的 T const 类型指针，如果 optional 没有值则返回 nullptr。

必须声明此重载以避免涉及隐式转换为 IBuffer 时出现的歧义。

模板参数：

**T** – 底层数据的类型。

参数:

optionalBufferPtr – 可能为空的 optional 对象。

返回：

指向常量 T 的指针，可能为 nullptr。

class ITensor : public virtual tensorrt_llm::runtime::IBuffer

公共类型

using UniquePtr = std::unique_ptr<ITensor>

using SharedPtr = std::shared_ptr<ITensor>

using UniqueConstPtr = std::unique_ptr<ITensor const>

using SharedConstPtr = std::shared_ptr<ITensor const>

using Shape = nvinfer1::Dims

using DimType64 = std::remove_reference_t<decltype(Shape::d[0])>

using TensorMap = runtime::StringPtrMap<runtime::ITensor>

公共函数

~ITensor() override = default

virtual Shape const &getShape() const = 0

返回张量维度。

template<SizeType32 n>
inline DimType64 getDimension() const

返回张量的第 n 个维度。如果 n 为负，则返回第 (nbDims - n) 个维度。TODO：迁移到 C++20 时替换为 constexpr 参数。

virtual void reshape(Shape const &dims) = 0

设置张量维度。张量的新大小将是 volume(dims)

inline virtual void resize(std::size_t newSize) override

调整缓冲区大小。如果新大小小于或等于当前容量，则此操作无效。

ITensor(ITensor const&) = delete

不允许复制。

ITensor &operator=(ITensor const&) = delete

不允许复制。

inline void squeeze(SizeType32 dim)

从此张量中移除给定的 *unit* 维度。

inline void unsqueeze(SizeType32 dim)

在指定位置添加一个 *unit* 维度。

inline bool shapeEquals(Shape const &other) const

inline bool shapeEquals(

std::initializer_list<SizeType32> const &other,

) const

template<typename T>
inline bool shapeEquals(

T const *dims,

SizeType32 count,

) const

公共静态函数

static inline std::int64_t volume(Shape const &dims)

返回维度的体积。如果 d.nbDims < 0 则返回 -1。

static inline std::size_t volumeNonNegative(Shape const &shape)

返回维度的体积。如果 d.nbDims < 0 则抛出异常。

static inline Shape strides(Shape const &dims)

返回 Shape 中每个维度的步长。

static Shape squeeze(Shape const &shape, SizeType32 dim)

从 shape 中移除给定的 *unit* 维度。

参数:

• shape – 要压缩（squeeze）的形状。

• dim – 应移除（“squeezed”）的维度。

返回：

一个新的没有单位维度的形状。

static Shape unsqueeze(Shape const &shape, SizeType32 dim)

在指定位置为 shape 添加一个 单位 维度。

参数:

• shape – 要取消压缩（unsqueeze）的形状。

• dim – 应该添加单位维度的维度。

返回：

添加了单位维度的新形状。

static UniquePtr slice(

SharedPtr tensor,

std::size_t offset,

std::size_t size,

)

在底层 tensor 上创建一个切片视图。该视图将具有与 tensor 相同的数据类型。

参数:

• tensor – 要查看的张量。

• offset – 视图相对于张量维度 0 的偏移量。

• size – 视图相对于张量维度 0 的大小。

返回：

buffer 的一个视图。

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::size_t offset,

std::size_t size,

)

static inline UniquePtr slice(SharedPtr tensor, std::size_t offset)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::size_t offset,

)

static UniquePtr slice(

SharedPtr tensor,

Shape const &offsetDims,

DimType64 size,

)

参数:

• offsetDims – 多维度的偏移量。

• tensor – 要查看的张量。

• offsetDims – 视图的偏移维度。

• size – 视图相对于 offsetDims 中最后一个维度的大小。

• offsetDims – 指定所有维度。

抛出：

Whenever – 偏移量溢出或最后一个维度的偏移量+大小溢出时。

返回：

形状为 [size, 其余维度] 或 [size] 的视图，当

static inline UniquePtr slice(

SharedPtr tensor,

std::initializer_list<DimType64> const &offsetDims,

DimType64 size,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

Shape const &offsetDims,

std::size_t size,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::initializer_list<DimType64> const &offsetDims,

std::size_t size,

)

static inline UniquePtr slice(

SharedPtr tensor,

Shape const &offsetDims,

)

当省略 size 时，返回最后一个维度上的其余切片。

static inline UniquePtr slice(

SharedPtr tensor,

std::initializer_list<DimType64> const &offsetDims,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

Shape const &offsetDims,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr slice(

TConstPtr &&tensor,

std::initializer_list<DimType64> const &offsetDims,

)

static inline UniquePtr at(SharedPtr tensor, Shape const &offsetDims)

参数:

offsetDims – 指定所有维度。

返回：

仅该点的块，形状为 [其余维度] 或 [1]，当

static inline UniquePtr at(

SharedPtr tensor,

std::initializer_list<DimType64> const &offsetDims,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr at(

TConstPtr &&tensor,

Shape const &offsetDims,

)

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline ITensor::UniqueConstPtr at(

TConstPtr &&tensor,

std::initializer_list<DimType64> const &offsetDims,

)

static UniquePtr view(IBuffer::SharedPtr buffer, Shape const &dims)

返回给定形状的底层 buffer（或张量）的视图。

参数:

• tensor – 要查看的张量。

• shape – 视图的形状。

返回：

tensor 的一个视图。

template<typename TConstPtr, std::enable_if_t<std::is_const_v<PointerElementType<TConstPtr>>, int> = 0>
static inline UniqueConstPtr view(

TConstPtr &&tensor,

Shape const &dims,

)

static inline UniquePtr view(SharedPtr tensor)

返回底层 tensor 的视图，该视图可以独立地重新整形。

参数:

tensor – 要查看的张量。

返回：

tensor 的一个视图。

static inline UniquePtr flattenN(

SharedPtr tensor,

std::int64_t sliceN = -1,

)

返回底层 tensor 的展平视图，该视图可以独立地重新整形。

参数:

• tensor – 要展平的张量。

• sliceN – 展平后切取前 N 个元素。-1 表示取整个展平的张量。

返回：

tensor 的展平视图。

static UniquePtr wrap(

void *data,

nvinfer1::DataType type,

Shape const &shape,

std::size_t capacity,

)

将给定的 data 封装到 ITensor 中。ITensor 不拥有底层 data，并且不能重新整形超出 capacity。

参数:

• data – 要封装的数据。

• type – data 的数据类型。

• shape – 张量的形状。

• capacity – 缓冲区的容量。

返回：

一个 ITensor。

static inline UniquePtr wrap(

void *data,

nvinfer1::DataType type,

Shape const &shape,

)

template<typename T>
static inline UniquePtr wrap(

T *data,

Shape const &shape,

std::size_t capacity,

)

template<typename T>
static inline UniquePtr wrap(

T *data,

Shape const &shape,

)

template<typename T>
static inline UniquePtr wrap(

std::vector<T> &v,

Shape const &shape,

)

static Shape makeShape(

std::initializer_list<DimType64> const &dims,

)

一个方便函数，用于创建具有给定维度的张量形状。

static std::string toString(Shape const &dims)

一个方便函数，用于将张量形状转换为 string。

static inline bool shapeEquals(Shape const &lhs, Shape const &rhs)

一个方便函数，用于比较形状。

template<typename T>
static inline bool shapeEquals(

Shape const &lhs,

T const *dims,

SizeType32 count,

)

一个方便函数，用于比较形状。

保护函数

ITensor() = default

受保护的静态函数

static inline DimType64 castSize(size_t newSize)

友元

friend class ITensorBindings

loraCache.h

namespace tensorrt_llm

namespace runtime

函数

std::string to_string(LoraCache::TaskLayerModuleConfig const &v)

std::ostream &operator<<(

std::ostream &os,

LoraCache::TaskLayerModuleConfig const &v,

)

class LoraCache

#include <loraCache.h>

LoraCache

使用 LRU 逐出策略缓存 LoRA 权重。

放入缓存的任务被标记为进行中，在标记为已完成之前无法被逐出。

一个缓存页存储一个最佳大小的 LoRA。一个页的大小为 [numSlots x pageWidth]。最佳大小的 LoRA 是指配置了 optimalAdapterSize 的 LoRA。

概念上，一个槽对应于一个 r=1、1 层、1 模块的输入/输出权重集。页宽度设置为最小模块中的权重数量。

每页的槽数量等于 ceilDiv(最佳大小 LoRA 中的权重数量, 最小模块中的权重数量)

缓存页在一个或多个块上分配

公共类型

using TensorPtr = ITensor::SharedPtr

using TaskIdType = std::uint64_t

using TaskLayerModuleConfigListPtr = std::shared_ptr<std::vector<TaskLayerModuleConfig>>

公共函数

LoraCache(

LoraCachePageManagerConfig const &pageManagerConfig,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

BufferManager const &bufferManager,

)

param[in] pageManagerConfig: 一个 LoraCachePageManagerConfig param[in] modelConfig: 一个 ModelConfig param[in] worldConfig: 一个 WorldConfig param[in] bufferManager: 一个 BufferManager 仅用于分配页块

void put(

TaskIdType taskId,

TensorPtr weights,

TensorPtr config,

bool load = true,

)

将任务放入缓存，为其申请页，并可选地加载任务权重。

参数:

• taskId – [in] 任务 ID

• weights – [in] LoRA 权重张量

• config – [in] LoRA 配置张量

• load – [in] 如果为 true，则在返回前加载权重，否则不加载

void loadWeights(

TaskIdType taskId,

TensorPtr weights,

TensorPtr config,

)

加载任务权重。此方法必须在 put 后调用。它被设计为在 put 返回 load = false 后异步调用

参数:

• taslId – [in] 任务 ID

• weights – [in] LoRA 权重张量

• config – [in] LoRA 配置张量

inline bool isLoaded(TaskIdType taskId) const

参数:

taskId – [in] 任务 ID

返回：

&#8212; 如果任务已加载（权重已就位）则为 true，否则为 false

bool isDone(TaskIdType taskId) const

参数:

taskId – [in] 任务 ID

返回：

&#8212; 如果任务已标记为完成且可以被逐出则为 true

inline bool has(TaskIdType taskId) const

参数:

taskId – [in] 任务 ID

返回：

&#8212; 如果任务在缓存中（不一定已加载）则为 true，否则为 false

std::vector<TaskLayerModuleConfig> const &get(TaskIdType taskId)

参数:

taskId – [in] 任务 ID

返回：

&#8212; Value 对象列表，包含指向任务权重的指针

void bump(TaskIdType taskId)

标记任务为最近使用

参数:

taskId – [in] 任务 ID

void markTaskDone(TaskIdType taskId)

将任务标记为完成，表示它可以被逐出

参数:

taskId – [in] 任务 ID

void markAllDone()

将缓存中的所有任务标记为完成

SizeType32 determineNumPages(TaskIdType taskId) const

参数:

taskId – [in] 任务 ID

返回：

&#8212; 存储给定任务所需的页数

SizeType32 determineNumPages(TensorPtr config) const

参数:

config – [in] LoRA 配置张量

返回：

&#8212; 存储由 config 张量配置的任务所需的页数

bool fits(TensorPtr config) const

参数:

config – [in] 一个 LoRA 配置张量

返回：

&#8212; 如果任务适合缓存，则为 true，否则为 false

void copyTask(

TaskIdType taskId,

LoraCache &deviceCache,

bool markDone = false,

)

将任务复制到另一个缓存。缓存必须具有相同的页大小。

参数:

• taskId – [in] 要复制的任务 ID

• otherCache – [in] 要将任务移动到的 LoraCache

• markDone – [in] 在复制时将已复制的任务标记为完成

SizeType32 getNumPages() const

返回：

&#8212; 分配给缓存的总页数（无论是否使用）

ITensor::SharedConstPtr getPagePtr(size_t pageId) const

参数:

pageId – [in] 页 ID

返回：

&#8212; 指向页的 const 指针

公共静态函数

static std::vector<LoraCache::TaskLayerModuleConfig> copyToPages(

TensorPtr weights,

TensorPtr config,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

std::unordered_map<SizeType32, LoraModule> moduleIdToModel,

BufferManager const &manager,

std::vector<TensorPtr> const &pages,

std::vector<std::size_t> const &pageIds,

)

将任务权重复制到缓存页。

参数:

• weights – [in] 任务权重

• config – [in] 任务配置张量

• modelConfig – [in] 一个 ModelConfig

• worldConfig – [in] 一个 WorldConfig

• modelIdToModel – [in] 从 LoRA 模块 ID 到 LoraModule 的映射

• manager – [in] 一个 BufferManager 用于执行复制操作的管理器

• pages – [out] 要将权重复制到的页张量列表

• pageIds – [in] 页的 ID 列表

返回：

&#8212; 缓存 Value 对象列表

static void splitTransposeCpu(

ITensor &output,

ITensor const &input,

SizeType32 tpSize,

SizeType32 tpRank,

)

将输入的第二个维度分割成 tpSize 部分，并将 tpRank 部分写入到 output

参数:

• output – [out] 输出张量

• input – [in] 输入张量

• tpSize – [in] 分割数量

• tpRank – [in] 要写入到 output 的分割部分

私有类型

enum ValueStatus

值

enumerator kVALUE_STATUS_MISSING

enumerator kVALUE_STATUS_PROCESSING

enumerator kVALUE_STATUS_LOADED

using TaskValuePtr = std::shared_ptr<TaskValue>

私有函数

void loadWeights(

TaskValue &cacheValue,

TensorPtr weights,

TensorPtr config,

)

void bumpTaskInProgress(TaskIdType taskId)

ValueStatus getStatus(TaskIdType taskId) const

std::vector<std::size_t> claimPagesWithEvict(SizeType32 numPages)

声明 numPages 页面，如果需要则驱逐任务

参数:

numPages – [in] 要声明的页面数量

抛出：

std::runtime_error – 如果无法声明所有页面

返回：

– 页面 ID 列表

std::map<size_t, std::pair<size_t, SizeType32>> copyTaskMapPages(

TaskValue &targetTaskValue,

TaskValue const &sourceTaskValue,

std::vector<size_t> const &targetPageIds,

LoraCache const &targetCache,

)

内部辅助方法，用于 copyTask 内部。本身不是线程安全的。

私有成员

LoraCachePageManagerConfig mPageManagerConfig

ModelConfig mModelConfig

WorldConfig mWorldConfig

mutable std::mutex mPagesMutex

std::unique_ptr<LoraCachePageManager> mCachePageManager

mutable std::mutex mCacheMutex

std::unordered_map<TaskIdType, TaskValuePtr> mCacheMap

std::list<TaskIdType> mInProgressTasks

std::list<TaskIdType> mDoneTasks

std::vector<std::unique_ptr<BufferManager>> mDeviceBufferManagers

std::unique_ptr<BufferManager> mBufferManager

std::unordered_map<SizeType32, LoraModule> mModuleIdToModule

私有静态函数

template<typename T>
static void splitTransposeCpuInner(

ITensor &output,

ITensor const &input,

SizeType32 tpSize,

SizeType32 tpRank,

)

struct TaskLayerModuleConfig

#include <loraCache.h>

包含单个层/模块的信息。这些配置的列表与每个任务相关联，可用于填充运行时张量。

公共函数

std::string toString() const

bool operator==(LoraCache::TaskLayerModuleConfig const &o) const

公共成员

std::size_t pageId

SizeType32 slotIdx

SizeType32 inSize

SizeType32 outSize

SizeType32 moduleId

SizeType32 layerId

SizeType32 adapterSize

SizeType32 numSlots

std::int64_t weightsInPointer

std::int64_t weightsOutPointer

std::optional<std::int64_t> scalingVecPointer

友元

friend class TaskLayerModuleConfigBindings

struct TaskValue

保存单个任务的配置和状态。

公共函数

TaskValue() = delete

~TaskValue() = default

inline TaskValue(

std::vector<std::size_t> const &pageIds,

TaskLayerModuleConfigListPtr const &configs,

std::list<TaskIdType>::iterator it,

bool inProgress,

bool loaded,

bool done,

bool loadInProgress = false,

)

inline TaskValue(TaskValue &&o) noexcept

inline TaskValue &operator=(TaskValue &&o)

公共成员

std::vector<std::size_t> pageIds

TaskLayerModuleConfigListPtr configs

std::list<TaskIdType>::iterator it

bool inProgress

bool loaded

bool done

将任务标记为已完成。这用于在加载期间将任务标记为已完成。如果在加载结束时（put、loadweights 或 copyTask 结束时）done=true，则该任务将被标记为已完成。

bool loadInProgress

表示权重正在通过 put 或 loadWeights 加载。这用于阻止同一任务的并发 loadWeights 调用。

class LoraCacheFullException : public tensorrt_llm::runtime::LoraExpectedException

公共函数

explicit LoraCacheFullException(std::string const &msg)

~LoraCacheFullException() noexcept override

class LoraCachePageManager

#include <loraCache.h>

持有 LoRA 缓存页的内存，并管理整页的分配和释放。内存预先分配在主机或设备上。

注意，此类非线程安全。

公共类型

using TensorPtr = ITensor::SharedPtr

公共函数

LoraCachePageManager(

使用给定的配置和 BufferManager 创建一个 LoraCachePageManager。,

BufferManager const &bufferManager,

)

参数:

• config – [in] 一个 LoraCachePageManagerConfig

• bufferManager – [in] 用于分配页块的 BufferManager

std::optional<std::vector<std::size_t>> claimPages(

SizeType32 numPages,

)

请求页

参数:

numPages – [in] 要声明的页面数量

返回：

一个元组，第一个值是布尔值，指示是否请求到页。如果第一个值为 true，则第二个值将包含页 ID 列表。

SizeType32 numAvailablePages() const

获取管理器中可用（空闲）页的数量

返回：

管理器中的空闲页数量

void releasePages(std::vector<std::size_t> const &pages)

释放给定的页

参数:

pages – [in] 要释放（空闲）的页列表

ITensor::SharedConstPtr blockPtr(SizeType32 blockIdx) const

返回给定页块的指针

参数:

blockIdx; – [in]

返回：

&#8212; 指向页块的指针

ITensor::SharedConstPtr pagePtr(std::size_t pageIdx) const

返回给定页的指针

参数:

pageIdx – [in]

返回：

&#8212; 指向页的 const 指针

ITensor::SharedPtr mutablePagePtr(std::size_t pageIdx)

返回给定页的指针

参数:

pageIdx – [in]

返回：

&#8212; 指向页的可变指针

私有函数

void initialize(BufferManager const &bufferManager)

私有成员

std::vector<TensorPtr> mPageBlocks

std::deque<std::size_t> mFreePageIds

std::vector<std::uint8_t> mIsPageFree

LoraCachePageManagerConfig const mConfig

class LoraExpectedException : public std::runtime_error

由 tensorrt_llm::runtime::LoraCacheFullException 继承

公共函数

explicit LoraExpectedException(std::string const &msg)

~LoraExpectedException() noexcept override

loraModule.h

namespace tensorrt_llm

namespace runtime

函数

inline std::ostream &operator<<(

std::ostream &output,

module – [in],

)

LoraModule const &module

公共类型

enum class ModuleType : SizeType32

值

enumerator kINVALID

enumerator kATTN_QKV

enumerator kATTN_Q

enumerator kATTN_K

enumerator kATTN_V

enumerator kATTN_DENSE

enumerator kMLP_H_TO_4H

enumerator kMLP_4H_TO_H

enumerator kMLP_GATE

enumerator kCROSS_ATTN_QKV

enumerator kCROSS_ATTN_Q

enumerator kCROSS_ATTN_K

enumerator kCROSS_ATTN_V

enumerator kCROSS_ATTN_DENSE

enumerator kMOE_H_TO_4H

enumerator kMOE_4H_TO_H

enumerator kMOE_GATE

enumerator kMOE_ROUTER

enumerator kMLP_ROUTER

enumerator kMLP_GATE_UP

using TensorPtr = ITensor::SharedPtr

公共函数

inline explicit constexpr LoraModule(

t – [in] 模块类型,

inDim – [in] 输入维度,

SizeType32 outDim,

bool inDimFirst,

bool outDimFirst,

SizeType32 inTpSplitDim,

SizeType32 outTpSplitDim,

) noexcept

inline explicit constexpr LoraModule() noexcept

explicit constexpr LoraModule(LoraModule const &o) = default

constexpr LoraModule &operator=(LoraModule const &o) = default

inline SizeType32 constexpr flattenedInOutSize(

SizeType32 adapterSize,

bool isDora,

) const noexcept

inline SizeType32 constexpr inSize(

SizeType32 adapterSize,

) const noexcept

inline SizeType32 constexpr outSize(

SizeType32 adapterSize,

) const noexcept

inline SizeType32 constexpr localInSize(

SizeType32 adapterSize,

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localOutSize(

SizeType32 adapterSize,

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localScalesSize(

SizeType32 tpSize,

bool isDora,

) const noexcept

inline SizeType32 constexpr localInDim(

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localOutDim(

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localInAdapterSize(

SizeType32 adapterSize,

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localOutAdapterSize(

SizeType32 adapterSize,

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localInOutSize(

SizeType32 adapterSize,

SizeType32 tpSize,

) const noexcept

inline SizeType32 constexpr localTotalSize(

SizeType32 adapterSize,

SizeType32 tpSize,

bool isDora,

) const noexcept

inline SizeType32 constexpr value() const noexcept

inline std::string_view constexpr name() const noexcept

inline SizeType32 constexpr inDim() const noexcept

inline SizeType32 constexpr outDim() const noexcept

inline bool constexpr inDimFirst() const noexcept

inline bool constexpr outDimFirst() const noexcept

inline SizeType32 constexpr inTpSplitDim() const noexcept

inline SizeType32 constexpr outTpSplitDim() const noexcept

公共静态函数

static std::vector<LoraModule> createLoraModules(

std::vector<std::string> const &loraModuleNames,

SizeType32 hiddenSize,

SizeType32 mlpHiddenSize,

SizeType32 numAttentionHeads,

SizeType32 numKvAttentionHeads,

SizeType32 attentionHeadSize,

SizeType32 tpSize,

SizeType32 numExperts,

)

static inline ModuleType constexpr toModuleType(

std::string_view const &name,

)

static inline std::string_view constexpr toModuleName(

ModuleType t,

) noexcept

static inline std::string_view constexpr toModuleName(SizeType32 id)

私有成员

ModuleType mType

SizeType32 mInDim

SizeType32 mOutDim

bool mInDimFirst

bool mOutDimFirst

SizeType32 mInTpSplitDim

SizeType32 mOutTpSplitDim

modelConfig.h

namespace tensorrt_llm

namespace runtime

class ModelConfig

公共类型

枚举类 class ModelVariant : std::int32_t

值

枚举值 kGpt

枚举值 kChatGlm

枚举值 kGlm

枚举值 kMamba

枚举值 kRecurrentGemma

枚举值 kEncDec

枚举类 class LayerType : std::int32_t

值

枚举值 kATTENTION

枚举值 kRECURRENT

枚举值 kLINEAR

枚举值 kNOOP

枚举类 class KVCacheType : std::int32_t

值

枚举值 kCONTINUOUS

枚举值 kPAGED

枚举值 kDISABLED

枚举类 class ManageWeightsType : std::int32_t

值

枚举值 kDisabled

枚举值 kEnabled

公共函数

inline explicit ModelConfig(

SizeType32 词汇量大小,

SizeType32 层数,

SizeType32 注意力层数,

SizeType32 循环层数,

SizeType32 注意力头数,

SizeType32 hiddenSize,

nvinfer1::DataType dtype,

)

inline SizeType32 constexpr 获取词汇量大小() const noexcept

inline SizeType32 constexpr 获取填充后的词汇量大小(

SizeType32 世界大小,

) const noexcept

inline SizeType32 计算本地层数(

LayerType 层类型,

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

) const

inline SizeType32 计算较低级别层数(

LayerType 层类型,

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

) const

inline SizeType32 获取层数(SizeType32 流水线并行度 = 1) const

inline SizeType32 获取注意力层数(

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

) const

inline SizeType32 获取循环层数(

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

) const

inline SizeType32 constexpr 获取注意力头数() const noexcept

inline SizeType32 获取KV头数(SizeType32 层索引) const

inline void 设置KV头数(SizeType32 KV头数)

inline void 设置交叉KV头数(SizeType32 KV头数)

inline SizeType32 constexpr 获取隐藏层大小() const noexcept

inline SizeType32 constexpr 获取编码器隐藏层大小() const noexcept

inline void constexpr 设置编码器隐藏层大小(

SizeType32 编码器隐藏层大小,

) noexcept

inline SizeType32 constexpr 获取每个头的大小() const noexcept

inline void constexpr 设置每个头的大小(SizeType32 每个头的大小) noexcept

inline nvinfer1::DataType constexpr 获取数据类型() const noexcept

inline bool constexpr 使用Gpt注意力插件() const noexcept

inline bool constexpr 使用GemmAllReduce插件() const noexcept

inline void constexpr 使用Gpt注意力插件(

bool 使用Gpt注意力插件,

) noexcept

inline void constexpr 使用GemmAllReduce插件(

bool 使用GemmAllReduce插件,

) noexcept

inline bool constexpr 使用MambaConv1d插件() const noexcept

inline void constexpr 使用MambaConv1d插件(

bool 使用MambaConv1d插件,

) noexcept

inline bool constexpr 使用打包输入() const noexcept

inline void constexpr 使用打包输入(bool 输入是否打包) noexcept

inline bool constexpr 使用分页状态() const noexcept

inline void constexpr 使用分页状态(bool 分页状态) noexcept

inline SizeType32 constexpr getTokensPerBlock() const noexcept

inline void constexpr setTokensPerBlock(

SizeType32 TokensPerBlock,

) noexcept

inline common::QuantMode constexpr getQuantMode() const noexcept

inline void constexpr setQuantMode(

common::QuantMode QuantMode,

) noexcept

inline bool constexpr supportsInflightBatching() const noexcept

inline SizeType32 constexpr getMaxBatchSize() const noexcept

inline void constexpr setMaxBatchSize(

SizeType32 maxBatchSize,

) noexcept

inline SizeType32 constexpr getMaxBeamWidth() const noexcept

inline void constexpr setMaxBeamWidth(

SizeType32 maxBeamWidth,

) noexcept

inline SizeType32 constexpr getMaxInputLen() const noexcept

inline void constexpr setMaxInputLen(SizeType32 maxInputLen) noexcept

inline SizeType32 constexpr getMaxSequenceLen() const noexcept

inline void constexpr setMaxSequenceLen(

SizeType32 maxSequenceLen,

) noexcept

inline std::optional<SizeType32> constexpr getMaxNumTokens(

) const noexcept

inline void constexpr setMaxNumTokens(

std::optional<SizeType32> maxNumTokens,

) noexcept

inline SizeType32 constexpr getMaxEncoderLen() const noexcept

inline void constexpr setMaxEncoderLen(

SizeType32 maxEncoderLen,

) noexcept

inline bool constexpr usePromptTuning() const noexcept

inline bool constexpr useMrope() const noexcept

inline void constexpr setUseMrope(bool useMrope) noexcept

inline SizeType32 constexpr getMaxPositionEmbeddings() const noexcept

inline void constexpr setMaxPositionEmbeddings(

SizeType32 maxPositionEmbeddings,

) noexcept

inline SizeType32 constexpr getRotaryEmbeddingDim() const noexcept

inline void constexpr setRotaryEmbeddingDim(

SizeType32 rotaryEmbeddingDim,

) noexcept

inline SizeType32 constexpr getMaxPromptEmbeddingTableSize(

) const noexcept

inline void constexpr setMaxPromptEmbeddingTableSize(

SizeType32 maxPromptEmbeddingTableSize,

) noexcept

inline bool constexpr computeContextLogits() const noexcept

inline void constexpr computeContextLogits(

bool computeContextLogits,

) noexcept

inline bool constexpr computeGenerationLogits() const noexcept

inline void constexpr computeGenerationLogits(

bool computeGenerationLogits,

) noexcept

inline ModelVariant getModelVariant() const

inline void setModelVariant(ModelVariant modelVariant)

inline SizeType32 getMaxDecodingDraftTokens() const

inline SizeType32 constexpr getMaxDecodingTokens() const noexcept

inline void constexpr setContextFMHA(bool contextFMHA) noexcept

inline bool constexpr getContextFMHA() const noexcept

inline void constexpr setPagedContextFMHA(

bool pagedContextFMHA,

) noexcept

inline bool constexpr getPagedContextFMHA() const noexcept

inline void constexpr setPpReduceScatter(

bool ppReduceScatter,

) noexcept

inline bool constexpr getPpReduceScatter() const noexcept

inline bool constexpr useLoraPlugin() const noexcept

inline void constexpr useLoraPlugin(bool useLoraPlugin) noexcept

inline std::vector<LoraModule> const &获取LoraModules(

) const noexcept

inline void 设置LoraModules(

std::vector<LoraModule> const &loraModules,

) noexcept

inline SizeType32 constexpr 获取MLP隐藏层大小() const noexcept

inline void constexpr 设置MLP隐藏层大小(

SizeType32 mlpHiddenSize,

) noexcept

inline bool constexpr 是否启用KV缓存() const noexcept

inline bool constexpr 是否分页KV缓存() const noexcept

inline bool constexpr 是否连续KV缓存() const noexcept

inline KVCacheType constexpr 获取KV缓存类型() const noexcept

inline void constexpr 设置KV缓存类型(

KVCacheType kvCacheType,

) noexcept

inline bool constexpr 使用交叉注意力() const noexcept

inline void constexpr 设置是否使用交叉注意力(

bool useCrossAttention,

) noexcept

inline bool constexpr 使用位置嵌入() const noexcept

inline void constexpr 设置是否使用位置嵌入(

bool usePositionEmbedding,

) noexcept

inline bool constexpr 使用TokenType嵌入() const noexcept

inline void constexpr 设置是否使用TokenType嵌入(

bool useTokenTypeEmbedding,

) noexcept

inline SizeType32 constexpr 获取最大Lora Rank() const noexcept

inline void constexpr 设置最大Lora Rank(SizeType32 maxLoraRank) noexcept

inline void 设置推测解码模式(

SpeculativeDecodingMode mode,

) noexcept

inline bool 是否具有推测解码模块() const noexcept

inline SpeculativeDecodingModule const &获取推测解码模块(

) const noexcept

inline std::shared_ptr<SpeculativeDecodingModule const> 获取推测解码模块指针(

) const noexcept

inline std::shared_ptr<SpeculativeDecodingModule> 获取推测解码模块指针(

) noexcept

inline void 设置推测解码模块(

std::shared_ptr<SpeculativeDecodingModule> const &speculativeDecodingModule,

) noexcept

inline void 重置推测解码模块() noexcept

inline void 启用无缝先行解码(

SizeType32 maxDraftTokens,

) noexcept

inline void 禁用无缝先行解码() noexcept

inline nvinfer1::DataType 获取Kv数据类型() const

inline bool constexpr 是否基于Transformer() const noexcept

inline bool 是否具有RNN配置() const noexcept

inline std::optional<RnnConfig> 获取RNN配置() const noexcept

inline void 设置RNN配置(RnnConfig const &rnnConfig) noexcept

inline bool constexpr 是否基于RNN() const noexcept

inline std::vector<LayerType> const &获取层类型(

) const noexcept

inline void 设置层类型(

std::vector<LayerType> const &layerTypes,

) noexcept

inline SpeculativeDecodingMode constexpr 获取推测解码模式(

) const noexcept

inline void 设置Logits数据类型(nvinfer1::DataType inputDtype) noexcept

inline nvinfer1::DataType constexpr 获取Logits数据类型() const noexcept

inline void 设置Gemm AllReduce数据类型(

nvinfer1::DataType inputDtype,

) noexcept

inline nvinfer1::DataType constexpr getGemmAllReduceDtype(

) const noexcept

inline void setUseShapeInference(bool useShapeInference) noexcept

inline bool useShapeInference() const noexcept

inline ManageWeightsType getManageWeightsType() const noexcept

inline void setManageWeightsType(

ManageWeightsType const manageWeightType,

) noexcept

inline std::string const &getModelName() const noexcept

inline void setModelName(std::string const &modelName)

inline std::vector<SizeType32> const &getNumKvHeadsPerLayer(

) const

inline std::pair<std::vector<SizeType32>::const_iterator, std::vector<SizeType32>::const_iterator> getNumKvHeadsPerLayerLocalRange(

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

bool isCrossAttention = false,

) const

inline void setNumKvHeadsPerLayer(

std::vector<SizeType32> const &headsPerLayer,

)

inline void setNumKvHeadsPerCrossLayer(

std::vector<SizeType32> const &headsPerLayer,

)

inline SizeType32 getSumLocalKvHeads(

SizeType32 流水线并行度 = 1,

SizeType32 流水线并行度级别 = 0,

bool isCrossAttention = false,

) const

inline bool constexpr skipCrossAttnBlocks() const noexcept

inline void constexpr setSkipCrossAttnBlocks(

bool skipCrossAttnBlocks,

) noexcept

inline std::optional<SizeType32> constexpr getNumLanguages(

) const noexcept

inline bool constexpr useLanguageAdapter() const noexcept

inline void constexpr setNumLanguages(

std::optional<SizeType32> numLanguages,

) noexcept

inline bool isMultiModal() const

inline bool isWhisper() const

公共静态函数

static inline KVCacheType KVCacheTypeFromString(std::string value)

static inline std::vector<SizeType32> getOptProfilesSplitPoints(

) noexcept

公共静态属性

static constexpr std::array kOPT_PROFILES_SPLIT_POINTS = {64, 128, 256, 512, 1024}

static constexpr SizeType32 kDEFAULT_NUM_TOKENS_PER_BLOCK = 64

私有成员

SizeType32 mVocabSize

SizeType32 mNbLayers

SizeType32 mNbAttentionLayers

SizeType32 mNbRnnLayers

SizeType32 mNbHeads

SizeType32 mHiddenSize

SizeType32 mSizePerHead

nvinfer1::DataType mDataType

bool mUseGptAttentionPlugin

bool mUseGemmAllReducePlugin

nvinfer1::DataType mGemmAllReduceDtype

bool mUseMambaConv1dPlugin

bool mInputPacked

bool mPagedState

SizeType32 mTokensPerBlock

common::QuantMode mQuantMode

SizeType32 mMaxBatchSize

SizeType32 mMaxBeamWidth

SizeType32 mMaxInputLen

SizeType32 mMaxSequenceLen

最大token数量

计算上下文logits

计算生成logits

模型变体

最大prompt嵌入表大小

使用Mrope

最大位置嵌入

旋转嵌入维度

上下文FMHA

分页上下文FMHA

PP Reduce Scatter

使用Lora插件

Lora模块

MLP隐藏层大小

最大Lora秩

RNN配置

KV缓存类型

最大编码器长度

编码器隐藏层大小

使用交叉注意力

使用位置嵌入

使用token类型嵌入

层类型

推测解码模块

推测解码模式

Logits数据类型

使用形状推断

管理权重类型

模型名称

每注意力层的KV头数量

每交叉注意力层的KV头数量

跳过交叉注意力块

语言数量

结构体 RnnConfig

公共成员

状态大小

卷积核

RNN隐藏层大小

RNN头大小

RNN卷积维度大小

lookaheadModule.h

namespace tensorrt_llm

namespace runtime

类 LookaheadModule

公共函数

LookaheadModule 构造函数

最大草稿路径长度,

最大解码草稿token数量,

) noexcept

LookaheadModule 构造函数

设置执行配置

配置,

)

获取执行配置

) const

私有成员

执行配置

request.h

namespace tensorrt_llm

namespace runtime

命名空间 decoder_batch

类 Request

公共类型

Tensor常量指针

Tensor指针

Buffer指针

公共函数

Request 构造函数

ID,

输入长度,

最大新token数量,

std::optional<SizeType32> endId = std::nullopt,

)

公共成员

TensorConstPtr ids

强制参数。

SizeType32 inputLen

std::optional<SizeType32> maxNewTokens

std::optional<SizeType32> endId

SizeType32 generatedTokensPerEngineStep

TensorPtr embeddingBias

TensorPtr badWordsList

TensorPtr stopWordsList

BufferPtr draftTokens

推测解码的可选参数。

std::optional<TensorPtr> draftLogits

TensorPtr medusaPaths

TensorPtr medusaTreeIds

nvinfer1::DataType dtype

std::optional<executor::LookaheadDecodingConfig> lookaheadRuntimeConfig

std::optional<executor::EagleConfig> eagleConfig

loraCachePageManagerConfig.h

namespace tensorrt_llm

namespace runtime

函数

inline std::ostream &operator<<(

std::ostream &os,

LoraCachePageManagerConfig const &c,

)

inline std::string to_string(LoraCachePageManagerConfig const &c)

class LoraCachePageManagerConfig

#include <loraCachePageManagerConfig.h>

LoraCachePageManager 的配置

有关页面、插槽和页块的描述，请参阅 LoraCache 文档。

公共函数

inline explicit constexpr LoraCachePageManagerConfig(

runtime::MemoryType memType,

nvinfer1::DataType dType,

SizeType32 totalNumPages,

SizeType32 maxPagesPerBlock,

SizeType32 slotsPerPage,

SizeType32 pageWidth,

SizeType32 numCopyStreams,

)

inline runtime::MemoryType constexpr getMemoryType() const noexcept

inline void constexpr setMemoryType(

runtime::MemoryType const &memoryType,

) noexcept

inline nvinfer1::DataType constexpr getDataType() const noexcept

inline void constexpr setDataType(

nvinfer1::DataType const &dtype,

) noexcept

inline SizeType32 constexpr getTotalNumPages() const noexcept

inline void constexpr setTotalNumPage(

SizeType32 const &totalNumPages,

) noexcept

inline SizeType32 constexpr getMaxPagesPerBlock() const noexcept

inline void constexpr setMaxPagesPerBlock(

SizeType32 const &maxPagesPerBlock,

) noexcept

inline SizeType32 constexpr getSlotsPerPage() const noexcept

inline void constexpr setSlotsPerPage(

SizeType32 const &slotsPerPage,

) noexcept

inline SizeType32 constexpr getPageWidth() const noexcept

inline void constexpr setPageWidth(

SizeType32 const &pageWidth,

) noexcept

inline bool constexpr getInitToZero() const noexcept

inline void constexpr setInitToZero(bool initToZero) noexcept

inline SizeType32 constexpr getNumCopyStreams() const noexcept

inline void constexpr setNumCopyStreams(

SizeType32 numCopyStreams,

) noexcept

私有成员

runtime::MemoryType mMemoryType

nvinfer1::DataType mDataType

SizeType32 总页数

SizeType32 每块最大页数

SizeType32 每页槽位

SizeType32 页宽度

SizeType32 复制流数量 = 1

bool 初始化为零

samplingConfig.h

宏定义

SET_FROM_OPTIONAL(varName, VarName, VarType)

namespace tensorrt_llm

namespace runtime

class 采样配置

公共函数

inline explicit 采样配置(SizeType32 束宽度 = 1)

inline explicit 采样配置(

std::vector<采样配置> const &configs,

)

inline explicit 采样配置(

executor::SamplingConfig const &samplingConfig,

std::optional<executor::ExternalDraftTokensConfig> const &externalDraftTokensConfig = std::nullopt,

)

inline bool 验证()

template<typename T>
inline bool 使用默认值(

OptVec<T> const &vec,

T defaultValue,

)

inline bool operator==(采样配置 const &other) const

inline SizeType32 获取返回 Beam 数量() const

inline SizeType32 按迭代次数获取 Beam 宽度(

int const decodingIter,

int const indexSC = 0,

) const noexcept

inline SizeType32 获取最大 Beam 宽度() const noexcept

公共成员

SizeType32 束宽度

std::optional<SizeType32> 返回序列数量

OptVec<FloatType> 温度

OptVec<FloatType> 原始温度

OptVec<SizeType32> 最小生成长度

OptVec<FloatType> 重复惩罚

OptVec<FloatType> 存在惩罚

OptVec<FloatType> 频率惩罚

OptVec<SizeType32> 禁止重复 N-gram 大小

OptVec<bool> 输出对数概率

OptVec<bool> 累计对数概率

OptVec<SizeType32> Top K 采样

OptVec<FloatType> Top P 采样

OptVec<uint64_t> 随机种子

OptVec<FloatType> Top P 衰减

OptVec<FloatType> 最小 Top P

OptVec<TokenIdType> Top P 重置 ID

OptVec<FloatType> 最小概率

OptVec<FloatType> Beam 搜索多样性率

OptVec<FloatType> 长度惩罚

OptVec<SizeType32> 提前停止

OptVec<std::vector<SizeType32>> 束宽度数组

OptVec<FloatType> 草稿接受阈值

OptVec<std::vector<SizeType32>> Medusa 头部的 Top K

std::optional<bool> 标准化对数概率

私有类型

using FloatType = float

template<typename T>
using OptVec = std::optional<std::vector<T>>

私有函数

template<typename T>
inline bool validateVec(

class GptJsonConfig,

OptVec<T> const &vec,

T min,

std::optional<T> max = std::nullopt,

)

私有静态函数

template<typename T>
static inline OptVec<T> fuseValues(

std::vector<SamplingConfig> const &configs,

std::function<OptVec<T>(size_t ci)> accessor,

T defaultValue,

)

medusaModule.h

namespace tensorrt_llm

namespace runtime

class MedusaModule : public tensorrt_llm::runtime::SpeculativeDecodingModule

公共类型

using TensorPtr = ITensor::SharedPtr

using MedusaChoices = std::vector<std::vector<SizeType32>>

公共函数

inline explicit MedusaModule(

SizeType32 maxAcceptedTokens,

SizeType32 maxDraftTokens,

) noexcept

inline explicit MedusaModule() noexcept

inline MedusaChoices const &getMedusaChoices() const noexcept

私有成员

MedusaChoices mDefaultMedusaChoices = {{0}, {0, 0}, {1}, {0, 1}, {2}, {0, 0, 0}, {1, 0}, {0, 2}, {3}, {0, 3}, {4}, {0, 4} , {2, 0} , {0, 5} , {0, 0, 1} , {5} , {0, 6} , {6} , {0, 7} , {0, 1, 0} , {1, 1} , {7} , {0, 8} , {0, 0, 2} , {3, 0} , {0, 9} , {8} , {9} , {1, 0, 0} , {0, 2, 0} , {1, 2} , {0, 0, 3} , {4, 0} , {2, 1} , {0, 0, 4} , {0, 0, 5} , {0, 0, 0, 0} , {0, 1, 1} , {0, 0, 6} , {0, 3, 0} , {5, 0} , {1, 3} , {0, 0, 7} , {0, 0, 8} , {0, 0, 9} , {6, 0} , {0, 4, 0} , {1, 4} , {7, 0} , {0, 1, 2} , {2, 0, 0} , {3, 1} , {2, 2} , {8, 0} , {0, 5, 0} , {1, 5} , {1, 0, 1} , {0, 2, 1} , {9, 0} , {0, 6, 0} , {0, 0, 0, 1} , {1, 6} , {0, 7, 0}}

memoryCounters.h

namespace tensorrt_llm

namespace runtime

class MemoryCounters

公共类型

using SizeType32 = std::size_t

using DiffType = std::ptrdiff_t

公共函数

MemoryCounters() = default

inline SizeType32 getGpu() const

inline SizeType32 getCpu() const

inline SizeType32 getPinned() const

inline SizeType32 getUVM() const

inline SizeType32 getPinnedPool() const

inline DiffType getGpuDiff() const

inline DiffType getCpuDiff() const

inline DiffType getPinnedDiff() const

inline DiffType getUVMDiff() const

inline DiffType getPinnedPoolDiff() const

template<MemoryType T>
inline void allocate(SizeType32 size)

void allocate(MemoryType memoryType, SizeType32 size)

template<MemoryType T>
inline void deallocate(SizeType32 size)

void deallocate(MemoryType memoryType, SizeType32 size)

std::string toString() const

公共静态函数

static MemoryCounters &getInstance()

static std::string bytesToString(SizeType32 bytes, int precision = 2)

static std::string bytesToString(DiffType bytes, int precision = 2)

私有成员

std::atomic<SizeType32> mGpu = {}

std::atomic<SizeType32> mCpu = {}

std::atomic<SizeType32> mPinned = {}

std::atomic<SizeType32> mUVM = {}

std::atomic<SizeType32> mPinnedPool = {}

std::atomic<DiffType> mGpuDiff = {}

std::atomic<DiffType> mCpuDiff = {}

std::atomic<DiffType> mPinnedDiff = {}

std::atomic<DiffType> mUVMDiff = {}

std::atomic<DiffType> mPinnedPoolDiff = {}

worldConfig.h

namespace tensorrt_llm

namespace runtime

class WorldConfig

公共函数

explicit WorldConfig(

SizeType32 tensorParallelism = 1,

SizeType32 流水线并行度 = 1,

SizeType32 contextParallelism = 1,

SizeType32 rank = 0,

SizeType32 gpusPerNode = kDefaultGpusPerNode,

std::optional<std::vector<SizeType32>> const &deviceIds = std::nullopt,

bool enableAttentionDP = false,

)

inline SizeType32 constexpr getSize() const noexcept

inline SizeType32 constexpr getTensorParallelism() const noexcept

inline bool constexpr isTensorParallel() const noexcept

inline SizeType32 constexpr getPipelineParallelism() const noexcept

inline bool constexpr isPipelineParallel() const noexcept

inline SizeType32 constexpr getContextParallelism() const noexcept

inline bool constexpr isContextParallel() const noexcept

inline SizeType32 constexpr getRank() const noexcept

inline SizeType32 constexpr getGpusPerNode() const noexcept

inline SizeType32 getGpusPerGroup() const noexcept

inline SizeType32 getDevice() const noexcept

inline SizeType32 getDeviceOf(SizeType32 rank) const noexcept

inline SizeType32 constexpr getPipelineParallelRank() const noexcept

inline SizeType32 constexpr getTensorParallelRank() const noexcept

inline SizeType32 constexpr getContextParallelRank() const noexcept

inline SizeType32 constexpr getLocalRank() const noexcept

inline SizeType32 constexpr getNodeRank() const noexcept

inline SizeType32 constexpr getNodeRankOf(

SizeType32 rank,

) const noexcept

inline bool constexpr isFirstPipelineParallelRank() const noexcept

inline bool constexpr isLastPipelineParallelRank() const noexcept

我的 rank 是否是其 pipeline 中的最后一个 rank？

inline bool constexpr isFirstTensorParallelRank() const noexcept

inline bool constexpr isFirstContextParallelRank() const noexcept

inline SizeType32 constexpr getLastRank() const noexcept

inline bool constexpr enableAttentionDP() const noexcept

std::vector<SizeType32> getPipelineParallelGroup() const

std::vector<SizeType32> getTensorParallelGroup() const

std::vector<SizeType32> getContextParallelGroup() const

bool validMpiConfig() const

公共静态函数

static WorldConfig mpi(

SizeType32 gpusPerNode = kDefaultGpusPerNode,

std::optional<SizeType32> tensorParallelism = std::nullopt,

std::optional<SizeType32> pipelineParallelism = std::nullopt,

std::optional<SizeType32> contextParallelism = std::nullopt,

std::optional<std::vector<SizeType32>> const &deviceIds = std::nullopt,

bool enableAttentionDP = false,

)

公共静态属性

static SizeType32 constexpr kDefaultGpusPerNode = 1

私有成员

SizeType32 mTensorParallelism

SizeType32 mPipelineParallelism

SizeType32 mContextParallelism

SizeType32 mRank

SizeType32 mGpusPerNode

bool mEnableAttentionDP

std::vector<SizeType32> mDeviceIds

promptTuningParams.h

namespace tensorrt_llm

namespace runtime

template<typename TTensor>
class GenericPromptTuningParams

公共类型

using TensorPtr = TTensor

using SizeType32 = tensorrt_llm::runtime::SizeType32

公共函数

inline explicit GenericPromptTuningParams(

TensorPtr embeddingTable = TensorPtr(),

TensorPtr tasks = TensorPtr(),

TensorPtr vocabSize = TensorPtr(),

)

公共成员

TensorPtr embeddingTable

TensorPtr tasks

TensorPtr vocabSize

std::vector<bool> promptTuningEnabled

class PromptTuningParams : public tensorrt_llm::runtime::GenericPromptTuningParams<ITensor::SharedPtr>

公共类型

using TensorPtr = ITensor::SharedPtr

using SizeType32 = GenericPromptTuningParams::SizeType32

公共函数

inline explicit PromptTuningParams(

TensorPtr embeddingTable = nullptr,

TensorPtr tasks = nullptr,

TensorPtr vocabSize = nullptr,

)

void fillTasksTensor(

TensorPtr tasksHost,

const SizeType32 batchSize,

const SizeType32 numContextRequests,

std::vector<SizeType32> const &reqBeamWidths,

std::vector<SizeType32> const &reqPromptLengths,

BufferManager const &manager,

bool packedInput,

)

rawEngine.h

namespace tensorrt_llm

namespace runtime

class RawEngine

公共类型

enum Type

值

enumerator FilePath

enumerator AddressWithSize

enumerator HostMemory

公共函数

inline explicit RawEngine(std::filesystem::path enginePath) noexcept

inline explicit RawEngine(

void const *engineAddr,

std::size_t engineSize,

) noexcept

inline explicit RawEngine(

nvinfer1::IHostMemory const *engineBuffer,

) noexcept

inline Type getType() const

inline std::filesystem::path getPath() const

inline std::optional<std::filesystem::path> getPathOpt() const

inline void setPath(std::filesystem::path enginePath)

inline std::optional<std::map<std::string, tensorrt_llm::executor::Tensor>> const &getManagedWeightsMapOpt(

) const

inline void setManagedWeightsMap(

std::map<std::string, tensorrt_llm::executor::Tensor> managedWeightsMap,

)

inline void const *getAddress() const

inline std::size_t getSize() const

inline nvinfer1::IHostMemory const *getHostMemory() const

公共成员

void const *mEngineAddr = {}

std::size_t mEngineSize = {}

私有成员

Type mType

std::optional<std::filesystem::path> mEnginePath

struct tensorrt_llm::runtime::RawEngine

nvinfer1::IHostMemory const *mEngineBuffer = {}

std::optional<std::map<std::string, tensorrt_llm::executor::Tensor>> mManagedWeightsMap

decoderState.h

namespace tensorrt_llm

namespace runtime

namespace decoder

class BeamSearchBuffers

公共函数

explicit BeamSearchBuffers(BufferManager const &bufferManager)

void reshape(SizeType32 maxBeamWidth, SizeType32 maxSequenceLength)

公共成员

DecodingOutput::BeamHypotheses mOutputBeamHypotheses

DecodingOutput::TensorPtr mCumLogProbsTmp

SizeType32 mNumSMs

class DecoderState

公共类型

using TensorPtr = ITensor::SharedPtr

using LlmRequestPtr = std::shared_ptr<tensorrt_llm::batch_manager::LlmRequest>

using RequestVector = std::vector<LlmRequestPtr>

using DecodingInputPtr = std::unique_ptr<DecodingInput>

using DecodingOutputPtr = std::unique_ptr<DecodingOutput>

公共函数

DecoderState(

nvinfer1::DataType dtype,

BufferManager const &bufferManager,

)

void allocateSpeculativeDecodingBuffers(

SpeculativeDecodingMode speculativeDecodingMode,

nvinfer1::DataType dtype,

BufferManager const &bufferManager,

)

为推测解码分配缓冲区。

void setup(

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

BufferManager const &bufferManager,

)

void setupSpeculativeDecoding(

SpeculativeDecodingMode const &speculativeDecodingMode,

SizeType32 maxTokensPerEngineStep,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

BufferManager const &bufferManager,

)

为推测解码设置缓冲区。

void setupExplicitDraftTokens(

ExplicitDraftTokensBuffers::Inputs explicitDraftTokensBuffers,

) const

为 ExplicitDraftTokens 解码设置缓冲区。

void setupLookahead(

LookaheadDecodingBuffers lookaheadDecodingBuffers,

) const

为 Lookahead 解码设置缓冲区。

void setupEagle(EagleBuffers::Inputs eagleBuffers) const

为 Eagle 解码设置缓冲区。

void disableLookahead(RequestVector const &genRequests)

禁用 Lookahead 解码。

TensorPtr getFinishedSum() const

返回：

[batchSize]，每个请求已完成序列的数量，位于 GPU 上

TensorPtr getFinishReasons() const

返回：

[batchSize, beamWidth]，FinishedState 值，位于 GPU 上

TensorPtr getIds() const

返回：

[batchSize, maxBeamWidth, maxInputLength + maxNewTokens]，包含输入 token id 和生成的 token id，无填充，位于 GPU 上。对于 Beam Search，包含未收集的数据。

TensorPtr getIds(SizeType32 batchIdx) const

参数:

batchIdx – 批次的索引

返回：

[maxBeamWidth, maxInputLength + maxNewTokens]，包含请求 batchIdx 的输入 token id 和生成的 token id，无填充，位于 GPU 上。对于 Beam Search，包含未收集的数据。

TensorPtr getGatheredIds() const

返回：

[batchSize, maxBeamWidth, maxInputLength + maxNewTokens]，仅用于 Beam Search。包含已收集的 token id，无填充，位于 GPU 上。

TensorPtr getGatheredIds(SizeType32 batchIdx) const

参数:

batchIdx – 批次的索引

返回：

[batchSize, maxBeamWidth, maxInputLength + maxNewTokens]，仅用于 beam search。它包含 request batchIdx 的无填充的 gathered token id，位于 gpu 上。

TensorPtr getParentIds() const

返回：

[batchSize, maxBeamWidth, maxInputLength + maxNewTokens]，包含 beam search 期间收集的无填充的 parent id，位于 gpu 上。

TensorPtr getCumLogProbs() const

返回：

[batchSize, maxBeamWidth]，累积对数概率（每个 beam），在 GPU 上

TensorPtr getCumLogProbs(SizeType32 batchIdx) const

返回：

[maxBeamWidth]，累积对数概率（每个 beam），位于 gpu 上。

TensorPtr getLogProbs() const

返回：

[batchSize, maxBeamWidth, maxSequenceLength]，对数概率（每个 beam），在 GPU 上

TensorPtr getLogProbs(SizeType32 batchIdx) const

返回：

[maxBeamWidth, maxSequenceLength]，对数概率（每个 beam），位于 gpu 上。

TensorPtr getSequenceLengths() const

返回：

[batchSize, maxBeamWidth]，序列长度，位于 gpu 上。

TensorPtr getAllNewTokens() const

获取上一次前向传播中生成的 maxTokensPerStep 个 token。

返回：

[maxTokensPerStep, batchSize, maxBeamWidth]，上一次前向传播中生成的 token，位于 gpu 上。

TensorPtr getNextDraftTokens() const

返回：

[batchSize, maxDraftTokens]，预测的用于下一步的 draft token，位于 gpu 上。

TensorPtr getPrevDraftTokensLengths() const

返回：

[batchSize]，预测的用于上一步的 draft token 长度，位于 gpu 上。

TensorPtr getNextDraftTokensLengths() const

返回：

[batchSize]，预测的用于下一步的 draft token 长度，位于 gpu 上。

TensorPtr getAcceptedLengthsCumSum() const

返回：

[batchSize + 1]，已接受的 draft token 长度的独占累积和，位于 gpu 上。

TensorPtr getAcceptedPackedPaths() const

返回：

[batchSize, maxAcceptedDraftTokensPerStep]，已接受的路径打包到连续张量中，位于 gpu 上。

TensorPtr getFinishedSteps() const

返回：

[maxTokensPerStep, batchSize, beamWidth]，类型为 FinishedState 的完成状态，位于 gpu 上。

SizeType32 getActualBatchSize() const

void setActualBatchSize(SizeType32 actualBatchSize)

SizeType32 getMaxBeamWidth() const

SizeType32 getMaxSequenceLength() const

SizeType32 getMaxDecodingDecoderTokens() const

SizeType32 getMaxDecodingEngineTokens() const

std::vector<SizeType32> const &getNumDecodingEngineTokens() const

获取批次中所有请求的 token 数量。

返回：

批次中所有请求的 token 数量。

SizeType32 getNumDecodingEngineTokens(SizeType32 batchIdx) const

获取批次中特定请求的 token 数量。

参数:

batchIdx – 请求在批次中的索引。

返回：

指定请求的 token 数量。

void setNumDecodingEngineTokens(

SizeType32 batchIdx,

SizeType32 numTokens,

)

设置批次中特定请求的 token 数量。

参数:

• batchIdx – 请求在批次中的索引。

• numTokens – 指定请求的 token 数量。

SpeculativeDecodingMode getSpeculativeDecodingMode() const

BeamSearchBuffers const &getBeamSearchBuffers() const

流式模式下 beam search 的工作空间。

DecodingInput &getJointDecodingInput() const

解码器的有状态输入。为 maxBatchSize 个槽位分配。

DecodingOutput &getJointDecodingOutput() const

解码器的有状态输出。为 maxBatchSize 个槽位分配。

私有成员

SizeType32 mActualBatchSize = {}

SizeType32 mMaxBatchSize = {}

SizeType32 mMaxBeamWidth = {}

SizeType32 mMaxSequenceLength = {}

DecodingInputPtr mJointDecodingInput

解码器的有状态输入。为 maxBatchSize 个槽位分配。

DecodingOutputPtr mJointDecodingOutput

解码器的有状态输出。为 maxBatchSize 个槽位分配。

TensorPtr mFinishedSteps

[maxTokensPerStep, batchSize, beamWidth] 类型为 FinishedState 的完成状态，对应 maxTokensPerStep 的每个生成 token，位于 gpu 上。

std::unique_ptr<BeamSearchBuffers> mBeamSearchBuffers

流式模式下 beam search 的工作空间。

SizeType32 mMaxDecodingDecoderTokens = {1}

SizeType32 mMaxDecodingEngineTokens = {1}

std::vector<SizeType32> mNumDecodingEngineTokens

[batchSize]，每个请求的令牌数量。

SpeculativeDecodingMode mSpeculativeDecodingMode = {SpeculativeDecodingMode::None()}

runtimeDefaults.h

namespace tensorrt_llm

namespace runtime

struct RuntimeDefaults

公共函数

inline RuntimeDefaults(

std::optional<std::vector<SizeType32>> maxAttentionWindowVec,

std::optional<SizeType32> sinkTokenLength,

)

RuntimeDefaults() = default

公共成员

std::optional<std::vector<SizeType32>> maxAttentionWindowVec

std::optional<SizeType32> sinkTokenLength

eagleModule.h

namespace tensorrt_llm

namespace runtime

class EagleModule : public tensorrt_llm::runtime::SpeculativeDecodingModule

公共函数

inline explicit EagleModule(

最大草稿路径长度,

最大解码草稿token数量,

SizeType32 numTransformersLayer,

SizeType32 maxNonLeafNodesPerLayer,

) noexcept

inline explicit EagleModule() noexcept

inline executor::EagleChoices const &getDefaultEagleChoices(

) const noexcept

inline SizeType32 getNumTransformerLayers() const noexcept

inline SizeType32 getMaxNonLeafNodesPerLayer() const noexcept

私有成员

SizeType32 mNumTransformersLayer

SizeType32 mMaxNonLeafNodesPerLayer

executor::EagleChoices mDefaultEagleChoices = {{0}, {0, 0}, {1}, {0, 1}, {2}, {0, 0, 0}, {1, 0}, {0, 2}, {3}, {0, 3}, {4}, {0, 4}, {2, 0}, {0, 5}, {0, 0, 1}, {5}, {0, 6}, {6} , {0, 7}, {0, 1, 0}, {1, 1}, {7} , {0, 8}, {0, 0, 2}, {3, 0}, {0, 9}, {8} , {9} , {1, 0, 0}, {0, 2, 0}, {1, 2} , {0, 0, 3}, {4, 0}, {2, 1} , {0, 0, 4}, {0, 0, 5} , {0, 0, 0, 0} , {0, 1, 1} , {0, 0, 6} , {0, 3, 0} , {5, 0} , {1, 3} , {0, 0, 7} , {0, 0, 8} , {0, 0, 9} , {6, 0} , {0, 4, 0} , {1, 4} , {7, 0} , {0, 1, 2} , {2, 0, 0} , {3, 1} , {2, 2} , {8, 0} , {0, 5, 0} , {1, 5} , {1, 0, 1} , {0, 2, 1} , {9, 0} , {0, 6, 0} , {0, 0, 0, 1} , {1, 6} , {0, 7, 0}}

speculativeDecodingMode.h

namespace tensorrt_llm

namespace runtime

class SpeculativeDecodingMode

公共类型

using UnderlyingType = std::uint8_t

公共函数

inline bool constexpr isNone() const

inline bool constexpr isDraftTokensExternal() const

inline bool constexpr isMedusa() const

inline bool constexpr isLookaheadDecoding() const

inline bool constexpr isExplicitDraftTokens() const

inline bool constexpr isEagle() const

inline bool constexpr updatesPositionIds() const

inline bool constexpr requiresAttentionMask() const

inline bool constexpr predictsDraftTokens() const

inline bool constexpr needsKVCacheRewind() const

inline bool constexpr variableDraftLength() const

inline bool constexpr hasDraftLogits() const

inline bool constexpr needsDecoderPrologue() const

inline bool operator==(SpeculativeDecodingMode const &other) const

inline explicit constexpr SpeculativeDecodingMode(

UnderlyingType state,

)

公共静态函数

static inline auto constexpr None()

static inline auto constexpr DraftTokensExternal()

static inline auto constexpr Medusa()

static inline auto constexpr LookaheadDecoding()

static inline auto constexpr ExplicitDraftTokens()

static inline auto constexpr Eagle()

私有函数

inline bool constexpr anyBitSet(UnderlyingType bits) const

inline bool constexpr allBitSet(UnderlyingType bits) const

私有成员

UnderlyingType mState = {kNone}

私有静态属性

static UnderlyingType constexpr kNone = {1U << 0U}

static UnderlyingType constexpr kDraftTokensExternal = {1U << 1U}

static UnderlyingType constexpr kMedusa = {1U << 2U}

static UnderlyingType constexpr kLookaheadDecoding = {1U << 3U}

static UnderlyingType constexpr kExplicitDraftTokens = {1U << 4U}

static UnderlyingType constexpr kEagle = {1U << 5U}

speculativeDecodingModule.h

namespace tensorrt_llm

namespace runtime

class SpeculativeDecodingModule

子类包括 tensorrt_llm::runtime::EagleModule, tensorrt_llm::runtime::LookaheadModule, tensorrt_llm::runtime::MedusaModule

公共函数

inline explicit SpeculativeDecodingModule(

最大草稿路径长度,

最大解码草稿token数量,

SizeType32 maxNumPaths,

) noexcept

inline explicit SpeculativeDecodingModule() noexcept

virtual ~SpeculativeDecodingModule() = default

SpeculativeDecodingModule(SpeculativeDecodingModule const &o) = default

SpeculativeDecodingModule &operator=(

SpeculativeDecodingModule const &o,

) = default

inline SizeType32 getMaxDraftPathLen() const noexcept

返回：

解码器一步可接受的最大草稿令牌数

inline SizeType32 getMaxPathLen() const noexcept

比草稿路径长度多一，用于从主头进行预测

返回：

一个请求在解码器一步中可增长的最大令牌数

inline SizeType32 getMaxDecodingDraftTokens() const noexcept

返回：

解码器一步处理的最大草稿令牌数

inline SizeType32 getMaxDecodingTokens() const noexcept

比解码草稿令牌多一，用于从主头进行预测

返回：

解码器一步处理的最大令牌数

inline SizeType32 getNumPackedMasks() const noexcept

inline SizeType32 getMaxNumPaths() const noexcept

inline void setMaxDraftTokens(SizeType32 maxDraftTokens) noexcept

inline void setMaxDraftPathLen(SizeType32 maxDraftPathLen) noexcept

inline void setMaxNumPaths(SizeType32 maxNumPaths) noexcept

私有函数

inline void computeNumPackedMasks() noexcept

私有成员

SizeType32 mMaxDraftPathLen

SizeType32 mMaxDecodingDraftTokens

SizeType32 mMaxNumPaths

SizeType32 mMaxNumPackedMasks

ipcNvlsMemory.h

namespace tensorrt_llm

namespace runtime

函数

void MPI_group_barrier(std::set<int> ranks)

bool ipcNvlsSupported()

IpcNvlsHandle *ipcNvlsAllocate(size_t size, std::set<int> ranks)

void ipcNvlsFree(IpcNvlsHandle *handle)

template<typename T>
class DeviceAllocationNvls

公共函数

DeviceAllocationNvls() = default

inline ~DeviceAllocationNvls()

inline void reset(size_t size, std::set<int> ranks)

inline T *getMulticastPointer() const

inline T *getUnicastPointer() const

inline T **getIpcUnicastPointers()

inline size_t getCapacity() const

inline void free()

私有成员

size_t _capacity = 0

IpcNvlsHandle *_handle

struct IpcNvlsHandle

公共成员

size_t size = 0

uintptr_t uc_ptr = 0

uintptr_t mc_ptr = 0

std::vector<uintptr_t> ipc_uc_ptrs

CUdeviceptr uc_va

CUdeviceptr mc_va

std::vector<CUdeviceptr> ipc_uc_vas

CUmemGenericAllocationHandle uc_handle

CUmemGenericAllocationHandle mc_handle

std::vector<CUmemGenericAllocationHandle> ipc_uc_handles

tllmLogger.h

namespace tensorrt_llm

namespace runtime

class TllmLogger : public nvinfer1::ILogger

公共函数

void log(

Severity severity,

nvinfer1::AsciiChar const *msg,

) noexcept override

Severity getLevel()

void setLevel(Severity level)

statefulGptDecoderBatched.h

namespace tensorrt_llm

namespace runtime

class StatefulGptDecoderBatched : public tensorrt_llm::runtime::IStatefulGptDecoder

#include <statefulGptDecoderBatched.h>

支持在途批处理和有状态接口的 GPT 解码器类。

公共类型

using CudaStreamPtr = std::shared_ptr<CudaStream>

using TensorPtr = ITensor::SharedPtr

公共函数

StatefulGptDecoderBatched(

CudaStreamPtr stream,

nvinfer1::DataType dtype,

)

~StatefulGptDecoderBatched() override

virtual void setup(

executor::DecodingMode const &mode,

SizeType32 maxBatchSize,

SizeType32 maxBeamWidth,

SizeType32 maxAttentionWindow,

SizeType32 sinkTokenLength,

SizeType32 maxSequenceLength,

nvinfer1::DataType dtype,

ModelConfig const &modelConfig,

WorldConfig const &worldConfig,

) override

using TensorPtr = std::shared_ptr<ITensor>

virtual void newBatch(

GenerationInput const &inputs,

GenerationOutput const &outputs,

SamplingConfig const &samplingConfig,

ModelConfig const &modelConfig,

) override

在调用 forward() 之前设置解码器，也调用 reshapeBuffers。

virtual void forwardAsync(

decoder::Output &output,

decoder::Input const &input,

) override

在不阻塞主机线程的情况下，为所有请求运行一步。

virtual void forwardSync() override

等待对 forwardAsync 的最后一次调用完成。

virtual TensorPtr getIds() const override

返回：

[batchSize, beamWidth, maxSequenceLength]，所有 token id，在 GPU 上

virtual TensorPtr getGatheredIds() const override

返回：

[batchSize, beamWidth, maxSequenceLength] 在 gatherTree 后的 token id

virtual void finalize(

SamplingConfig const &samplingConfig,

) const override

收集所有请求的最终束搜索结果。

virtual TensorPtr getCumLogProbs() const override

返回：

[batchSize, maxBeamWidth]，累积对数概率（每个 beam），在 GPU 上

virtual TensorPtr getLogProbs() const override

返回：

[batchSize, maxBeamWidth, maxSequenceLength]，对数概率（每个 beam），在 GPU 上

virtual TensorPtr getNewTokens(SizeType32 iter = 0) const override

获取上一次 forward 传递中一步生成的 token。

参数:

**iter** – 要获取 token 的迭代次数，范围在 [0; maxTokensPerStep) 内

返回：

[batchSize, beamWidth]，在 iter 中生成的 token（每个 beam），在 GPU 上

virtual TensorPtr getNbFinished() const override

返回：

[1]，已完成序列的数量，在 pinned 主机内存中

私有成员

std::unique_ptr<GptDecoderBatched> mDecoder

CudaEvent mDecoderFinishEvent

CudaEvent mForwardEvent

TensorPtr mFinishedSum

TensorPtr mBatchSlotsSetup

TensorPtr mBatchSlotsDecoder