Runtime#

gptJsonConfig.h#

namespace tensorrt_llm
namespace runtime
class GptJsonConfig#

Public Functions

inline GptJsonConfig(
std::string name,
std::string version,
std::string precision,
SizeType32 tensorParallelism,
SizeType32 pipelineParallelism,
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#
inline std::optional<RuntimeDefaults> getRuntimeDefaults() const#
std::string engineFilename(
WorldConfig const &worldConfig,
std::string const &model,
) const#
inline std::string engineFilename(
WorldConfig const &worldConfig,
) const#

Public Static Functions

static GptJsonConfig parse(std::string const &json)#
static GptJsonConfig parse(std::istream &json)#
static GptJsonConfig parse(std::filesystem::path const &path)#

Private Members

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#
std::optional<RuntimeDefaults> mRuntimeDefaults#

worldConfig.h#

namespace tensorrt_llm
namespace runtime
class WorldConfig#

Public Functions

explicit WorldConfig(
SizeType32 tensorParallelism = 1,
SizeType32 pipelineParallelism = 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#

Is my rank the last rank in its pipeline?

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#

Public Static Functions

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,
)#

Public Static Attributes

static SizeType32 constexpr kDefaultGpusPerNode = 1#

Private Members

SizeType32 mTensorParallelism#
SizeType32 mPipelineParallelism#
SizeType32 mContextParallelism#
SizeType32 mRank#
SizeType32 mGpusPerNode#
bool mEnableAttentionDP#
std::vector<SizeType32> mDeviceIds#

common.h#

Defines

FMT_DIM#
namespace tensorrt_llm
namespace runtime

Typedefs

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>>#

Enums

enum class RequestType : std::int32_t#

Values:

enumerator kCONTEXT#
enumerator kGENERATION#
struct UniqueToken#

Public Functions

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

Public Members

TokenIdType tokenId#
TokenExtraIdType tokenExtraId#

ipcUtils.h#

namespace tensorrt_llm
namespace runtime

Functions

void lamportInitializeAll(
void *buffer_0,
void *buffer_1,
void *buffer_2,
size_t size,
)#
bool canAccessPeer(WorldConfig const &worldConfig)#
class AllReduceBuffers#

Public Types

using TensorPtr = ITensor::SharedPtr#

Public Functions

AllReduceBuffers(
SizeType32 maxBatchSize,
SizeType32 maxBeamWidth,
SizeType32 maxSequenceLength,
SizeType32 hiddenSize,
BufferManager const &manager,
WorldConfig const &worldConfig,
bool const fakeBuffers = false,
)#

Public Members

TensorPtr mAllReduceCommPtrs#
TensorPtr mFlagPtrs#
std::vector<runtime::IpcMemory> mIpcMemoryHandles#
class IpcMemory#

Public Types

using BufferPtr = IBuffer::SharedPtr#

Public Functions

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#

Public Static Attributes

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

Private Functions

void allocateIpcMemory(
std::size_t bufferSize,
BufferManager const &manager,
WorldConfig const &worldConfig,
)#
void destroyIpcMemory()#

Private Members

SizeType32 mTpRank#
std::vector<void*> mCommPtrs#
BufferPtr mBuffer#
bool mOpenIpc#

speculativeDecodingMode.h#

namespace tensorrt_llm
namespace runtime
class SpeculativeDecodingMode#

Public Types

using UnderlyingType = std::uint8_t#

Public Functions

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,
)#

Public Static Functions

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()#

Private Functions

inline bool constexpr anyBitSet(UnderlyingType bits) const#
inline bool constexpr allBitSet(UnderlyingType bits) const#

Private Members

UnderlyingType mState = {kNone}#

Private Static Attributes

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}#

memoryCounters.h#

namespace tensorrt_llm
namespace runtime
class MemoryCounters#

Public Types

using SizeType32 = std::size_t#
using DiffType = std::ptrdiff_t#

Public Functions

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#

Public Static Functions

static MemoryCounters &getInstance()#
static std::string bytesToString(SizeType32 bytes, int precision = 2)#
static std::string bytesToString(DiffType bytes, int precision = 2)#

Private Members

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 = {}#

runtimeDefaults.h#

namespace tensorrt_llm
namespace runtime
struct RuntimeDefaults#

Public Functions

inline RuntimeDefaults(
std::optional<std::vector<SizeType32>> maxAttentionWindowVec,
std::optional<SizeType32> sinkTokenLength,
)#
RuntimeDefaults() = default#

Public Members

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

virtualMemory.h#

namespace tensorrt_llm
namespace runtime

Functions

CudaVirtualMemoryManager &getVirtualMemoryManager()#
CudaVirtualMemoryAllocator getVirtualMemoryAllocator()#
void pushVirtualMemoryAllocator(
std::string const &tag,
CudaVirtualMemoryAllocator::RestoreMode mode,
std::shared_ptr<CudaStream> backStream,
)#
void popVirtualMemoryAllocator()#
class CudaVirtualMemoryAllocator#

Public Types

enum RestoreMode#

Values:

enumerator NONE#
enumerator MEMSET#
enumerator CPU#
enumerator PINNED#

Public Functions

inline explicit CudaVirtualMemoryAllocator(
std::shared_ptr<Configuration> config,
)#
inline explicit operator bool() const noexcept#
void allocate(Pointer *ptr, std::size_t n, int device) const#
void deallocate(Pointer ptr, std::size_t n) const#

Private Types

using CudaStreamPtr = std::shared_ptr<CudaStream>#
using Pointer = void*#

Private Members

std::shared_ptr<Configuration> mConfig#
class Configuration#

Public Functions

inline Configuration(
CudaVirtualMemoryManager &manager,
std::string tag,
RestoreMode mode,
CudaStreamPtr backStream,
)#

CudaVirtualMemoryAllocator::Configuration

Parameters:

  • manager – Manager used to track and manage virtual memories

  • tag – The tag for allocated memories

  • mode – Backed storage mode

  • backStream – The CUDA stream used for restoring memory content Note: Virtual Address Allocation is not async. The stream is not used in allocation.

inline std::size_t aligned(std::size_t n, int device = 0)#

Public Static Attributes

static Configuration backgroundConfiguration#

Private Functions

inline Configuration(
CudaVirtualMemoryManager &manager,
std::string tag,
RestoreMode mode,
CudaStreamPtr backStream,
bool background,
)#

Private Members

CudaVirtualMemoryManager &mManager#
std::string mTag#
CudaStreamPtr mBackStream#
std::atomic<std::size_t> mAlignment#
RestoreMode mMode#
bool mBackground = {}#

Friends

friend class CudaVirtualMemoryAllocator
friend void pushVirtualMemoryAllocator(
std::string const &tag,
RestoreMode mode,
std::shared_ptr<CudaStream> backStream,
)#
class CUDAVirtualMemoryChunk#
#include <virtualMemory.h>

CUDAVirtualMemoryChunk is a handle to a piece of CUDA memory allocation, providing the ability to release and rematerialize the allocation.

Public Types

enum Status#

Values:

enumerator INVALID#
enumerator RELEASED#
enumerator MATERIALIZED#
enumerator ERRORED#
using CreatorPtr = std::unique_ptr<Creator>#
using ConfiguratorPtr = std::unique_ptr<Configurator>#
using Configurators = std::vector<ConfiguratorPtr>#

Public Functions

inline Status status() const noexcept#
void materialize()#

Materialize this CUDAVirtualMemoryChunk. Shall be called only when status() == RELEASED.

Calls creator.create(), and then configurator.setup() for each configurator in order.

Stop at the first thrown exception and propagates it.

inline void release()#

Release this CUDAVirtualMemoryChunk. Shall be called only when status() == MATERIALIZED, or materialize() throws. Will be called automatically by destructor if necessary.

Calls configurator.teardown() for each configurator that setup() succeed in materialize() in reversed order, and then creator.release().

Never stops early upon exception. The last thrown exception will be propagated, and others logged.

CUDAVirtualMemoryChunk(CUDAVirtualMemoryChunk const&) = delete#
CUDAVirtualMemoryChunk &operator=(
CUDAVirtualMemoryChunk const&,
) = delete#
inline CUDAVirtualMemoryChunk(
CUDAVirtualMemoryChunk &&other,
) noexcept#
inline CUDAVirtualMemoryChunk &operator=(
CUDAVirtualMemoryChunk &&other,
)#
CUDAVirtualMemoryChunk() noexcept = default#
inline CUDAVirtualMemoryChunk(
CreatorPtr &&creator,
Configurators &&configurators,
)#
inline virtual ~CUDAVirtualMemoryChunk()#
inline explicit operator bool() const noexcept#

Test if this CUDAVirtualMemoryChunk is managing a memory block.

Private Functions

void _release(bool destructing)#

Private Members

size_t mState = 0#
CUmemGenericAllocationHandle mHandle = {}#
std::unique_ptr<Creator> mCreator#
std::vector<std::unique_ptr<Configurator>> mConfigurators#

Private Static Attributes

static constexpr size_t INVALID_STATE = static_cast<size_t>(-1)#
struct Configurator#
#include <virtualMemory.h>

CUDAVirtualMemoryChunk::Configurator is the interface to configure a CUmemGenericAllocationHandle:

  • Map into virtual address

  • Bind to multicast object

  • Backup and restore memory content

Subclassed by tensorrt_llm::runtime::MemsetConfigurator, tensorrt_llm::runtime::MulticastConfigurator, tensorrt_llm::runtime::OffloadConfigurator, tensorrt_llm::runtime::UnicastConfigurator

Public Functions

Configurator() = default#
virtual ~Configurator() = default#
Configurator(Configurator const&) = default#
Configurator &operator=(Configurator const&) = default#
Configurator(Configurator&&) = default#
Configurator &operator=(Configurator&&) = default#
virtual void setup(CUmemGenericAllocationHandle handle) = 0#
virtual void teardown(
CUmemGenericAllocationHandle handle,
bool destructing,
) = 0#
struct Creator#
#include <virtualMemory.h>

CUDAVirtualMemoryChunk::Creator is the interface to obtain a CUmemGenericAllocationHandle, either by creating one locally, or importing one from remote.

Subclassed by tensorrt_llm::runtime::LocalCreator< count >

Public Functions

Creator() = default#
virtual ~Creator() = default#
Creator(Creator const&) = default#
Creator &operator=(Creator const&) = default#
Creator(Creator&&) = default#
Creator &operator=(Creator&&) = default#
virtual CUmemGenericAllocationHandle create() = 0#
virtual void release(
CUmemGenericAllocationHandle handle,
bool destructing,
) = 0#
class CudaVirtualMemoryManager#

Public Functions

void add(
uintptr_t handle,
std::string tag,
CUDAVirtualMemoryChunk &&memory,
)#

Add memory to be managed by this manager.

The memory and internal state will remain valid if any exception is thrown.

Parameters:

  • handle – Unique handle provided to reference this memory in remove.

  • tag – Tag the memory, so this memory can be targeted in releaseWithTag and materializeWithTag.

  • memory – The CUDAVirtualMemory object.

void add(
uintptr_t handle,
std::string tag,
CUDAVirtualMemoryChunk::CreatorPtr &&creator,
CUDAVirtualMemoryChunk::Configurators &&configurators,
)#

Creates and adds memory to be managed by this manager. The created memory is automatically materialized.

The internal state will remain valid if any exception is thrown.

Parameters:

  • handle – Unique handle provided to reference this memory in remove.

  • tag – Tag the memory, so this memory can be targeted in releaseWithTag and materializeWithTag.

  • creator – The creator for the memory.

  • configurators – The configurators for the memory.

template<typename ...Configurators>
inline void add(
uintptr_t handle,
std::string tag,
CUDAVirtualMemoryChunk::CreatorPtr &&creator,
Configurators&&... configurators,
)#
CUDAVirtualMemoryChunk remove(uintptr_t handle) noexcept#

Remove the memory from the manager.

Parameters:

handle – The handle provided to add.

Returns:

The CUDAVirtualMemory object. If the handle is unknown, an empty CUDAVirtualMemory will be returned.

size_t releaseWithTag(std::string const &tag)#

Call release for CUDAVirtualMemoryChunk

objects with a given tag.

This function will always call

CUDAVirtualMemoryChunk::release on all selected objects. The last exception thrown by CUDAVirtualMemoryChunk::release will be rethrown, and others will be logged.

If any CUDAVirtualMemoryChunk threw an exception during release, it will be removed from the manager. Call retrieveBadHandles to retrieve handles of all CUDAVirtualMemoryChunk that got removed due to exception.

Parameters:

tag – the tag to select target memories.

Returns:

Number of objects selected.

size_t materializeWithTag(std::string const &tag)#

Call materialize for CUDAVirtualMemoryChunk

objects with a given tag.

This function will stop at the first

CUDAVirtualMemoryChunk::materialize that throws exception, and attempt to roll back previous successful materialize by calling release. The exception thrown by CUDAVirtualMemoryChunk::materialize will be rethrown, and any exception thrown by release will be logged.

If any CUDAVirtualMemoryChunk threw an exception during materialize or release, it will be removed from the manager. Successfully roll backed CUDAVirtualMemoryChunk will not be removed. Call retrieveBadHandles to retrieve handles of all CUDAVirtualMemoryChunk that got removed due to exception.

Parameters:

tag – the tag to select target memories.

Returns:

Number of objects selected.

std::vector<uintptr_t> retrieveBadHandles() noexcept#

Retrieve handles of all CUDAVirtualMemoryChunk that got removed due to exception and reset the list. The returned list may not include all removed CUDAVirtualMemoryChunk handles if OOM happened. This method is only for diagnostic purpose, and should not be called concurrently with other methods.

Returns:

The handle list.

Private Types

using PointerMemoryMap = std::map<uintptr_t, Entry>#
using TagEntryMap = std::multimap<std::string, PointerMemoryMap::iterator>#

Private Functions

CUDAVirtualMemoryChunk unsafeRemove(uintptr_t handle) noexcept#
void addBadHandle(uintptr_t handle) noexcept#

Private Members

std::mutex mMutex#
PointerMemoryMap mMemories#
TagEntryMap mEntries#
std::vector<uintptr_t> mBadHandles#
friend VirtualMemoryManagerTest
struct Entry#

Public Members

CUDAVirtualMemoryChunk mMemory#
TagEntryMap::iterator mEntryIt#
template<bool count = true>
struct LocalCreator : public tensorrt_llm::runtime::CUDAVirtualMemoryChunk::Creator#
#include <virtualMemory.h>

LocalCreator creates memory allocation locally through cuMemCreate.

Public Functions

inline LocalCreator(CUmemAllocationProp const &prop, size_t size)#
inline virtual CUmemGenericAllocationHandle create() override#
inline virtual void release(
CUmemGenericAllocationHandle handle,
bool destructing,
) override#

Public Members

CUmemAllocationProp mProp = {}#
size_t mSize = {}#
struct MemsetConfigurator : public tensorrt_llm::runtime::CUDAVirtualMemoryChunk::Configurator#
#include <virtualMemory.h>

MemsetConfigurator fills the memory with given value.

Public Functions

inline MemsetConfigurator(
CUdeviceptr address,
size_t size,
uint8_t value,
CUstream stream,
)#
inline virtual void setup(CUmemGenericAllocationHandle) override#
inline virtual void teardown(
CUmemGenericAllocationHandle,
bool,
) noexcept override#

Public Members

CUdeviceptr mAddress#
size_t mSize#
CUstream mStream = {}#
uint8_t mValue#
bool mFirstTime = true#
struct MulticastConfigurator : public tensorrt_llm::runtime::CUDAVirtualMemoryChunk::Configurator#
#include <virtualMemory.h>

MulticastConfigurator binds the allocation handle to the given multicast object and offset.

Public Functions

inline virtual void setup(
CUmemGenericAllocationHandle handle,
) override#
inline virtual void teardown(
CUmemGenericAllocationHandle,
bool,
) override#

Public Members

CUmemGenericAllocationHandle mMulticast#
size_t mBindOffset#
CUdevice mDevice#
size_t mSize#
struct OffloadConfigurator : public tensorrt_llm::runtime::CUDAVirtualMemoryChunk::Configurator#
#include <virtualMemory.h>

OffloadConfigurator offload the content of the allocation to the backup storage when teardown, and restore the content on the following setup.

Public Functions

inline OffloadConfigurator(
CUdeviceptr address,
size_t size,
MemoryType backType,
CUstream stream,
bool ondemand = false,
)#
virtual void setup(CUmemGenericAllocationHandle handle) override#
virtual void teardown(
CUmemGenericAllocationHandle handle,
bool destructing,
) override#

Public Members

CUdeviceptr mAddress#
size_t mSize#
MemoryType mBackType#
CUstream mStream#
bool mOndemand#
IBuffer::UniquePtr mBackedStorage#
struct UnicastConfigurator : public tensorrt_llm::runtime::CUDAVirtualMemoryChunk::Configurator#
#include <virtualMemory.h>

UnicastConfigurator maps the allocation handle into the specified unicast address range.

Public Functions

inline UnicastConfigurator(
CUdeviceptr address,
size_t size,
CUmemAccessDesc const &desc,
)#
inline virtual void setup(
CUmemGenericAllocationHandle handle,
) override#
inline virtual void teardown(
CUmemGenericAllocationHandle,
bool,
) override#

Public Members

CUdeviceptr mAddress#
size_t mSize#
CUmemAccessDesc mDesc#

bufferManager.h#

namespace tensorrt_llm
namespace runtime
class BufferManager#
#include <bufferManager.h>

A helper class for managing memory on host and device.

Public Types

using IBufferPtr = IBuffer::UniquePtr#
using ITensorPtr = ITensor::UniquePtr#
using CudaStreamPtr = std::shared_ptr<CudaStream>#
using CudaMemPoolPtr = std::shared_ptr<CudaMemPool>#

Public Functions

explicit BufferManager(CudaStreamPtr stream, bool trimPool = false)#

Construct a BufferManager.

Parameters:

cudaStream – [in] The cuda stream to use for all operations on GPU (allocation, de-allocation, copying, etc.).

inline ~BufferManager()#

Destructor.

IBufferPtr gpu(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Allocates an IBuffer of the given size on the GPU, using cudaMallocAsync.

ITensorPtr gpu(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Allocates an ITensor of the given dimensions on the GPU, using cudaMallocAsync.

IBufferPtr allocate(
MemoryType memoryType,
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Allocates an IBuffer of the given size and memory type.

ITensorPtr allocate(
MemoryType memoryType,
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Allocates an ITensor of the given dimensions and memory type.

inline IBufferPtr emptyBuffer(
MemoryType memoryType,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Create an empty IBuffer of the given memory type. It may be resized later.

inline ITensorPtr emptyTensor(
MemoryType memoryType,
tensorrt_llm::DataType type = kBYTE_TYPE,
) const#

Create an empty ITensor of the given memory type. It may be reshaped later.

void setMem(IBuffer &buffer, int32_t value) const#

Set the contents of the given buffer to value.

void setZero(IBuffer &buffer) const#

Set the contents of the given buffer to zero.

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

Copy src to dst.

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

Copy src to dst.

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

Copy src to dst.

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

Copy src to dst.

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

Copy src to dst.

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

Copy src into a new IBuffer with a potentially different memory type.

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

Copy src into a new ITensor with a potentially different memory type.

template<typename T>
inline IBufferPtr copyFrom(
std::vector<T> const &src,
MemoryType memoryType,
) const#

Copy src into a new IBuffer with a potentially different memory type.

template<typename T>
inline ITensorPtr copyFrom(
T *src,
tensorrt_llm::Dims dims,
MemoryType memoryType,
) const#

Copy src into a new ITensor with a potentially different memory type.

template<typename T>
inline ITensorPtr copyFrom(
std::vector<T> const &src,
tensorrt_llm::Dims dims,
MemoryType memoryType,
) const#

Copy src into a new ITensor with a potentially different memory type.

CudaStream const &getStream() const#

Get the underlying cuda stream.

std::size_t memoryPoolReserved() const#

The current size of the memory reserved by the memory pool.

std::size_t memoryPoolUsed() const#

The current size of the memory used by the memory pool.

std::size_t memoryPoolFree() const#

The current size of the memory free in the memory pool.

void memoryPoolTrimTo(std::size_t size)#

Try to trim the memory reserved by the pool to size bytes. This synchronizes implicitly with the stream.

Public Static Functions

static IBufferPtr gpuSync(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an IBuffer of the given size on the GPU, using cudaMalloc.

static ITensorPtr gpuSync(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an ITensor of the given dimensions on the GPU, using cudaMalloc.

static IBufferPtr cpu(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an IBuffer of the given size on the CPU.

static ITensorPtr cpu(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an ITensor of the given dimensions on the CPU.

static IBufferPtr pinned(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates a pinned IBuffer of the given size on the CPU.

static ITensorPtr pinned(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates a pinned ITensor of the given dimensions on the CPU.

static IBufferPtr pinnedPool(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates a pinned IBuffer of the given size on the CPU in the default memory pool.

static ITensorPtr pinnedPool(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates a pinned ITensor of the given dimensions on the CPU in the default memory pool.

static IBufferPtr managed(
std::size_t size,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an IBuffer of the given size in UVM.

static ITensorPtr managed(
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type = kBYTE_TYPE,
)#

Allocates an ITensor of the given dimensions in UVM.

static ITensorPtr ipcNvls(
std::set<int> ranks,
tensorrt_llm::Dims dims,
tensorrt_llm::DataType type,
)#

Allocates an ITensor of the given dimensions for NVLS.

Public Static Attributes

static auto constexpr kBYTE_TYPE = tensorrt_llm::DataType::kUINT8#

Private Members

CudaStreamPtr mStream#
CudaMemPoolPtr mPool#
bool const mTrimPool#

Friends

friend class ::BufferManagerTest

loraModule.h#

namespace tensorrt_llm
namespace runtime

Functions

inline std::ostream &operator<<(
std::ostream &output,
LoraModule const &module,
)#
class LoraModule#

Public Types

enum class ModuleType : SizeType32#

Values:

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#
enumerator kSHARED_EXPERT_H_TO_4H#
enumerator kSHARED_EXPERT_4H_TO_H#
enumerator kSHARED_EXPERT_GATE#
enumerator kMAMBA_IN_PROJ#
enumerator kMAMBA_OUT_PROJ#
enumerator kMOE_LATENT_FC1#
enumerator kMOE_LATENT_FC2#
using TensorPtr = ITensor::SharedPtr#

Public Functions

inline explicit constexpr LoraModule(
ModuleType const &t,
SizeType32 inDim,
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#

Public Static Functions

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,
SizeType32 sharedExpertHiddenSize = 0,
SizeType32 moeHiddenSize = 0,
SizeType32 mambaInProjSize = 0,
SizeType32 mambaInnerSize = 0,
SizeType32 moeLatentSize = 0,
)#
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)#

Private Members

ModuleType mType#
SizeType32 mInDim#
SizeType32 mOutDim#
bool mInDimFirst#
bool mOutDimFirst#
SizeType32 mInTpSplitDim#
SizeType32 mOutTpSplitDim#

cudaStream.h#

namespace tensorrt_llm
namespace runtime
class CudaStream#

Public Functions

inline explicit CudaStream(
unsigned int flags = cudaStreamNonBlocking,
int priority = 0,
)#

Creates a new cuda stream on the current device. The stream will be destroyed in the destructor.

Parameters:

  • flags – Flags for stream creation. See cudaStreamCreateWithFlags for a list of valid flags that can be passed.

  • priority – Priority of the stream. Lower numbers represent higher priorities. See cudaDeviceGetStreamPriorityRange for more information about the meaningful stream priorities that can be passed.

inline explicit CudaStream(
cudaStream_t stream,
int device,
bool ownsStream = true,
)#

Pass an existing cuda stream to this object.

Parameters:

  • stream – The stream to pass to this object.

  • device – The device on which the stream was created.

  • ownsStream – Whether this object owns the stream and destroys it in the destructor.

inline explicit CudaStream(cudaStream_t stream)#

Construct with an existing cuda stream or the default stream by passing nullptr.

inline int getDevice() const#

Returns the device on which the stream was created.

inline cudaStream_t get() const#

Returns the stream associated with this object.

inline void synchronize() const#

Synchronizes the stream.

inline void record(CudaEvent::pointer event) const#

Record an event on the stream.

inline void record(CudaEvent const &event) const#

Record an event on the stream.

inline void wait(CudaEvent::pointer event) const#

Wait for an event.

inline void wait(CudaEvent const &event) const#

Wait for an event.

Private Types

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

Private Members

StreamPtr mStream#
int mDevice = {-1}#
class Deleter#

Public Functions

inline explicit Deleter(bool ownsStream)#
inline explicit Deleter()#
inline constexpr void operator()(cudaStream_t stream) const#

Private Members

bool mOwnsStream#

cudaEvent.h#

namespace tensorrt_llm
namespace runtime
class CudaEvent#

Public Types

using pointer = cudaEvent_t#

Public Functions

inline explicit CudaEvent(unsigned int flags = cudaEventDisableTiming)#

Creates a new cuda event. The event will be destroyed in the destructor.

Parameters:

flags – Flags for event creation. By default, event timing is disabled.

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

Pass an existing cuda event to this object.

Parameters:

  • event – The event to pass to this object.

  • ownsEvent – Whether this object owns the event and destroys it in the destructor.

inline pointer get() const#

Returns the event associated with this object.

inline void synchronize() const#

Synchronizes the event.

Private Types

using element_type = std::remove_pointer_t<pointer>#
using EventPtr = std::unique_ptr<element_type, Deleter>#

Private Members

EventPtr mEvent#
class Deleter#

Public Functions

inline explicit Deleter(bool ownsEvent)#
inline explicit Deleter()#
inline constexpr void operator()(pointer event) const#

Private Members

bool mOwnsEvent#

modelConfig.h#

namespace tensorrt_llm
namespace runtime
class ModelConfig#

Public Types

enum class ModelVariant : std::int32_t#

Values:

enumerator kGpt#
enumerator kChatGlm#
enumerator kGlm#
enumerator kMamba#
enumerator kRecurrentGemma#
enumerator kEncDec#
enum class LayerType : std::int32_t#

Values:

enumerator kATTENTION#
enumerator kRECURRENT#
enumerator kLINEAR#
enumerator kNOOP#
enum class KVCacheType : std::int32_t#

Values:

enumerator kCONTINUOUS#
enumerator kPAGED#
enumerator kDISABLED#
enum class ManageWeightsType : std::int32_t#

Values:

enumerator kDisabled#
enumerator kEnabled#

Public Functions

inline explicit ModelConfig(
SizeType32 vocabSize,
SizeType32 nbLayers,
SizeType32 nbAttentionLayers,
SizeType32 nbRnnLayers,
SizeType32 nbHeads,
SizeType32 hiddenSize,
tensorrt_llm::DataType dtype,
)#
inline SizeType32 constexpr getVocabSize() const noexcept#
inline SizeType32 constexpr getVocabSizePadded(
SizeType32 worldSize,
) const noexcept#
inline SizeType32 countLocalLayers(
LayerType layerType,
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getFirstLocalLayer(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 countLowerRankLayers(
LayerType layerType,
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getNbLayers(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getNbAttentionLayers(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getNbRnnLayers(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getFirstLoraLayer(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline SizeType32 getNbLoraLayers(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
) const#
inline void setNbLoraLayers(SizeType32 nbLoraLayers)#
inline SizeType32 constexpr getNbHeads() const noexcept#
inline SizeType32 getNbKvHeads(SizeType32 layerIdx) const#
inline void setNbKvHeads(SizeType32 nbKvHeads)#
inline void setNbCrossKvHeads(SizeType32 nbKvHeads)#
inline SizeType32 constexpr getHiddenSize() const noexcept#
inline SizeType32 constexpr getEncoderHiddenSize() const noexcept#
inline void constexpr setEncoderHiddenSize(
SizeType32 encoderHiddenSize,
) noexcept#
inline SizeType32 constexpr getSizePerHead() const noexcept#
inline void constexpr setSizePerHead(SizeType32 sizePerHead) noexcept#
inline tensorrt_llm::DataType constexpr getDataType() const noexcept#
inline bool constexpr useGptAttentionPlugin() const noexcept#
inline bool constexpr useGemmAllReducePlugin() const noexcept#
inline void constexpr useGptAttentionPlugin(
bool useGptAttentionPlugin,
) noexcept#
inline void constexpr useGemmAllReducePlugin(
bool useGemmAllReducePlugin,
) noexcept#
inline bool constexpr useMambaConv1dPlugin() const noexcept#
inline void constexpr useMambaConv1dPlugin(
bool useMambaConv1dPlugin,
) noexcept#
inline bool constexpr usePackedInput() const noexcept#
inline void constexpr usePackedInput(bool inputPacked) noexcept#
inline bool constexpr usePagedState() const noexcept#
inline void constexpr usePagedState(bool pagedState) 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 &getLoraModules(
) const noexcept#
inline void setLoraModules(
std::vector<LoraModule> const &loraModules,
) noexcept#
inline SizeType32 constexpr getMlpHiddenSize() const noexcept#
inline void constexpr setMlpHiddenSize(
SizeType32 mlpHiddenSize,
) noexcept#
inline bool constexpr isKVCacheEnabled() const noexcept#
inline bool constexpr isPagedKVCache() const noexcept#
inline bool constexpr isContinuousKVCache() const noexcept#
inline KVCacheType constexpr getKVCacheType() const noexcept#
inline void constexpr setKVCacheType(
KVCacheType kvCacheType,
) noexcept#
inline bool constexpr useCrossAttention() const noexcept#
inline void constexpr setUseCrossAttention(
bool useCrossAttention,
) noexcept#
inline bool constexpr usePositionEmbedding() const noexcept#
inline void constexpr setUsePositionEmbedding(
bool usePositionEmbedding,
) noexcept#
inline bool constexpr useTokenTypeEmbedding() const noexcept#
inline void constexpr setUseTokenTypeEmbedding(
bool useTokenTypeEmbedding,
) noexcept#
inline SizeType32 constexpr getMaxLoraRank() const noexcept#
inline void constexpr setMaxLoraRank(SizeType32 maxLoraRank) noexcept#
inline void setSpeculativeDecodingMode(
SpeculativeDecodingMode mode,
) noexcept#
inline bool hasSpeculativeDecodingModule() const noexcept#
inline SpeculativeDecodingModule const &getSpeculativeDecodingModule(
) const noexcept#
inline std::shared_ptr<SpeculativeDecodingModule const> getSpeculativeDecodingModulePtr(
) const noexcept#
inline std::shared_ptr<SpeculativeDecodingModule> getSpeculativeDecodingModulePtr(
) noexcept#
inline void setSpeculativeDecodingModule(
std::shared_ptr<SpeculativeDecodingModule> const &speculativeDecodingModule,
) noexcept#
inline tensorrt_llm::DataType getKvDataType() const#
inline bool constexpr isTransformerBased() const noexcept#
inline bool hasRnnConfig() const noexcept#
inline std::optional<RnnConfig> getRnnConfig() const noexcept#
inline void setRnnConfig(RnnConfig const &rnnConfig) noexcept#
inline bool constexpr isRnnBased() const noexcept#
inline std::vector<LayerType> const &getLayerTypes(
) const noexcept#
inline void setLayerTypes(
std::vector<LayerType> const &layerTypes,
) noexcept#
inline SpeculativeDecodingMode constexpr getSpeculativeDecodingMode(
) const noexcept#
inline void setLogitsDtype(
tensorrt_llm::DataType inputDtype,
) noexcept#
inline tensorrt_llm::DataType constexpr getLogitsDtype(
) const noexcept#
inline void setGemmAllReduceDtype(
tensorrt_llm::DataType inputDtype,
) noexcept#
inline tensorrt_llm::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::vector<SizeType32> getNumKvHeadsForGivenLayers(
std::vector<SizeType32> const &layers,
bool isCrossAttention,
) const#
inline std::pair<std::vector<SizeType32>::const_iterator, std::vector<SizeType32>::const_iterator> getNumKvHeadsPerLayerLocalRange(
SizeType32 pipelineParallelism = 1,
SizeType32 pipelineParallelismRank = 0,
bool isCrossAttention = false,
) const#
inline void setNumKvHeadsPerLayer(
std::vector<SizeType32> const &headsPerLayer,
)#
inline void setNumKvHeadsPerCrossLayer(
std::vector<SizeType32> const &headsPerLayer,
)#
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#

Public Static Functions

static inline KVCacheType KVCacheTypeFromString(std::string value)#
static inline std::vector<SizeType32> getOptProfilesSplitPoints(
) noexcept#

Public Static Attributes

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

Private Members

SizeType32 mVocabSize#
SizeType32 mNbLayers#
SizeType32 mNbAttentionLayers#
SizeType32 mNbRnnLayers#
SizeType32 mNbHeads#
SizeType32 mHiddenSize#
SizeType32 mSizePerHead#
tensorrt_llm::DataType mDataType#
bool mUseGptAttentionPlugin#
bool mUseGemmAllReducePlugin#
tensorrt_llm::DataType mGemmAllReduceDtype#
bool mUseMambaConv1dPlugin#
bool mInputPacked#
bool mPagedState#
SizeType32 mTokensPerBlock#
common::QuantMode mQuantMode#
SizeType32 mMaxBatchSize#
SizeType32 mMaxBeamWidth#
SizeType32 mMaxInputLen#
SizeType32 mMaxSequenceLen#
std::optional<SizeType32> mMaxNumTokens#
bool mComputeContextLogits#
bool mComputeGenerationLogits#
ModelVariant mModelVariant#
SizeType32 mMaxPromptEmbeddingTableSize#
bool mUseMrope#
SizeType32 mMaxPositionEmbeddings#
SizeType32 mRotaryEmbeddingDim#
bool mContextFMHA#
bool mPagedContextFMHA#
bool mPpReduceScatter#
bool mUseLoraPlugin#
std::vector<LoraModule> mLoraModules#
SizeType32 mMlpHiddenSize#
SizeType32 mMaxLoraRank#
SizeType32 mNbLoraLayers = {0}#
std::optional<RnnConfig> mRnnConfig#
KVCacheType mKVCacheType = KVCacheType::kCONTINUOUS#
SizeType32 mMaxEncoderLen = {}#
SizeType32 mEncoderHiddenSize = {}#
bool mUseCrossAttention#
bool mUsePositionEmbedding#
bool mUseTokenTypeEmbedding#
std::vector<LayerType> mLayerTypes#
std::shared_ptr<SpeculativeDecodingModule> mSpeculativeDecodingModule#
SpeculativeDecodingMode mSpeculativeDecodingMode#
tensorrt_llm::DataType mLogitsDtype#
bool mUseShapeInference#
ManageWeightsType mManageWeightsType#
std::string mModelName#
std::vector<SizeType32> mNumKvHeadsPerAttentionLayer#
std::vector<SizeType32> mNumKvHeadsPerCrossAttentionLayer#
bool mSkipCrossAttnBlocks#
std::optional<SizeType32> mNumLanguages#
struct RnnConfig#

Public Members

SizeType32 stateSize = 0#
SizeType32 convKernel = 0#
SizeType32 rnnHiddenSize = 0#
SizeType32 rnnHeadSize = 0#
SizeType32 rnnConvDimSize = 0#

ipcNvlsMemory.h#

namespace tensorrt_llm
namespace runtime

Functions

void MPI_group_barrier(std::set<int> ranks)#
bool ipcNvlsSupported()#

Whether NVLS (NVLink SHARP) multicast memory can be allocated on this node. Checks only the static capability (CUDA driver >= 12010 and CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED). This is the precondition for ipcNvlsAllocate(): the allocator itself selects a fabric or POSIX-FD handle via getMemHandleType(), so single-node NVLS works over POSIX-FD even when the fabric/IMEX plane is not provisioned. The result is cached.

bool ipcNvlsFabricUsable()#

Whether the NVLink fabric/IMEX plane is provisioned so that NVLS multicast memory can actually be bound (not merely statically supported). Extends ipcNvlsSupported() with a live fabric probe (getMemHandleType() must resolve to CU_MEM_HANDLE_TYPE_FABRIC). Use this to decide whether NCCL may attempt NVLS: on an unprovisioned node the static multicast attribute is a false positive and NCCL aborts during init, so callers should disable NCCL_NVLS when this returns false. The (heavy) result is cached.

IpcNvlsHandle *ipcNvlsAllocate(size_t size, std::set<int> ranks)#
void ipcNvlsFree(IpcNvlsHandle *handle)#
template<typename T>
class DeviceAllocationNvls#

Public Functions

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()#

Private Members

size_t _capacity = 0#
IpcNvlsHandle *_handle#
struct IpcNvlsHandle#

Public Members

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#

iTensor.h#

namespace tensorrt_llm
namespace runtime

Functions

inline std::ostream &operator<<(
std::ostream &output,
ITensor::Shape const &dims,
)#

Utility function to print a shape.

std::ostream &operator<<(
std::ostream &output,
ITensor const &tensor,
)#

Utility function to print a tensor with its shape.

template<typename T>
T const *bufferCastOrNull(
ITensor::SharedConstPtr const &tensorPtr,
)#

Retrieves a T const typed pointer to the underlying data of the tensor pointed to by the tensorPtr, or nullptr if the tensorPtr is null.

This overload has to be declared to avoid ambiguity when an implicit conversion to IBuffer is involved.

Template Parameters:

T – The type of the underlying data.

Parameters:

tensorPtr – A possibly null shared ptr.

Returns:

A pointer to T const, possibly nullptr.

template<typename T>
T *bufferCastOrNull(
ITensor::SharedPtr const &tensorPtr,
)#

Retrieves a T typed pointer to the underlying data of the buffer pointed to by the tensorPtr, or nullptr if the tensorPtr is null.

This overload has to be declared to avoid ambiguity when an implicit conversion to IBuffer is involved.

Template Parameters:

T – The type of the underlying data.

Parameters:

tensorPtr – A possibly null shared ptr.

Returns:

A pointer to T, possibly nullptr.

template<typename T>
T *bufferCastOrNull(
std::optional<ITensor::SharedPtr> const &optionalTensorPtr,
)#

Retrieves a T typed pointer to the underlying data of the tensor pointed to by the tensor pointer contained in the optionalBufferPtr, or nullptr if the optional doesn’t have a value.

This overload has to be declared to avoid ambiguity when an implicit conversion to IBuffer is involved.

Template Parameters:

T – The type of the underlying data.

Parameters:

optionalBufferPtr – A possibly empty optional.

Returns:

A pointer to T, possibly nullptr.

template<typename T>
T const *bufferCastOrNull(
std::optional<ITensor::SharedConstPtr> const &optionalTensorPtr,
)#

Retrieves a T const typed pointer to the underlying data of the tensor pointed to by the tensor pointer contained in the optionalBufferPtr, or nullptr if the optional doesn’t have a value.

This overload has to be declared to avoid ambiguity when an implicit conversion to IBuffer is involved.

Template Parameters:

T – The type of the underlying data.

Parameters:

optionalBufferPtr – A possibly empty optional.

Returns:

A pointer to const T, possibly nullptr.

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

Public Types

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 = tensorrt_llm::Dims#
using DimType64 = std::remove_reference_t<decltype(Shape::d[0])>#
using TensorMap = runtime::StringPtrMap<runtime::ITensor>#

Public Functions

~ITensor() override = default#
virtual Shape const &getShape() const = 0#

Returns the tensor dimensions.

template<SizeType32 n>
inline DimType64 getDimension() const#

Returns the tensor n-th dimension. If n is negative, returns the (nbDims - n)th dimension.

virtual void reshape(Shape const &dims) = 0#

Sets the tensor dimensions. The new size of the tensor will be volume(dims)

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

Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity.

ITensor(ITensor const&) = delete#

Not allowed to copy.

ITensor &operator=(ITensor const&) = delete#

Not allowed to copy.

inline void squeeze(SizeType32 dim)#

Removes the given unit dimensions from this tensor.

inline void unsqueeze(SizeType32 dim)#

Adds a unit dimension at the specified position.

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#

Public Static Functions

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

Returns the volume of the dimensions. Returns -1 if d.nbDims < 0.

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

Returns the volume of the dimensions. Throws if d.nbDims < 0.

static inline Shape strides(Shape const &dims)#

Returns the strides of each dimemsion in a Shape.

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

Removes the given unit dimension from shape.

Parameters:

  • shape – The shape to squeeze.

  • dim – The dimension that should be removed (“squeezed”).

Returns:

A new shape without the unit dimension.

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

Add a unit dimension to shape at the specified position.

Parameters:

  • shape – The shape to unsqueeze.

  • dim – The dimension where unit dimension should be added.

Returns:

A new shape with the added unit dimension.

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

Creates a sliced view on the underlying tensor. The view will have the same data type as tensor.

Parameters:

  • tensor – The tensor to view.

  • offset – The offset of the view w.r.t. dimension 0 of the tensor.

  • size – The size of the view w.r.t. dimension 0 of the tensor.

Returns:

A view on the 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,
)#
Parameters:

  • offsetDims – The offset in multiple dimensions.

  • tensor – The tensor to view.

  • offsetDims – The offset dimensions of the view.

  • size – The size of the view w.r.t. the last dimension in offsetDims.

  • offsetDims – specifies all dimensions.

Throws:

Whenever – offset overflows or the last dimension offset+size overflows.

Returns:

A view of shape [size, the rest dimensions] or [size] when

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,
)#

return the rest slices at the last dimension when size omitted.

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)#
Parameters:

offsetDims – specifies all dimensions.

Returns:

Just the block at the point, with shape of [the rest dimensions] or [1] when

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)#

Returns a view on the underlying buffer (or tensor) with the given shape.

Parameters:

  • tensor – The tensor to view.

  • shape – The shape of the view.

Returns:

A view on the 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)#

Returns a view on the underlying tensor which can be independently reshaped.

Parameters:

tensor – The tensor to view.

Returns:

A view on the tensor.

static inline UniquePtr flattenN(
SharedPtr tensor,
std::int64_t sliceN = -1,
)#

Returns a flattened view on the underlying tensor which can be independently reshaped.

Parameters:

  • tensor – The tensor to flatten.

  • sliceN – Slice the first N elements after flattening. -1 means take the whole flattened tensor.

Returns:

A flatten view on the tensor.

static UniquePtr wrap(
void *data,
tensorrt_llm::DataType type,
Shape const &shape,
std::size_t capacity,
)#

Wraps the given data in an ITensor. The ITensor will not own the underlying data and cannot be reshaped beyond capacity.

Parameters:

  • data – The data to wrap.

  • type – The data type of the data.

  • shape – The shape of the tensor.

  • capacity – The capacity of the buffer.

Returns:

An ITensor.

static inline UniquePtr wrap(
void *data,
tensorrt_llm::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,
)#

A convenience function to create a tensor shape with the given dimensions.

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

A convenience function for converting a tensor shape to a string.

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

A convenience function to compare shapes.

template<typename T>
static inline bool shapeEquals(
Shape const &lhs,
T const *dims,
SizeType32 count,
)#

A convenience function to compare shapes.

Protected Functions

ITensor() = default#

Protected Static Functions

static inline DimType64 castSize(size_t newSize)#

Friends

friend class ITensorBindings

loraCache.h#

namespace tensorrt_llm
namespace runtime

Functions

std::string to_string(LoraCache::TaskLayerModuleConfig const &v)#
std::ostream &operator<<(
std::ostream &os,
LoraCache::TaskLayerModuleConfig const &v,
)#
class LoraCache#
#include <loraCache.h>

LoraCache

Caches LoRA weights with LRU eviction policy.

Tasks put in the cache are marked in progress and can not be evicted, until they are marked done.

A cache page holds a optimally sized LoRA. A page is of size [numSlots x pageWidth] An optimally size LoRA is on that has the configured optimalAdapterSize.

Conceptually a slot corresponds to a r=1, 1-layer, 1-module set of in/out weights. Page width is set to the number of weights in smallest module.

The number of slots per page is then ceilDiv(num weights in optimally sized LoRA, num weights in smallest module)

Cache pages are allocated on one or more blocks

Public Types

using TensorPtr = ITensor::SharedPtr#
using TaskIdType = std::uint64_t#
using TaskLayerModuleConfigListPtr = std::shared_ptr<std::vector<TaskLayerModuleConfig>>#

Public Functions

LoraCache(
LoraCachePageManagerConfig const &pageManagerConfig,
ModelConfig const &modelConfig,
WorldConfig const &worldConfig,
BufferManager const &bufferManager,
)#

param[in] pageManagerConfig: a LoraCachePageManagerConfig param[in] modelConfig: a ModelConfig param[in] worldConfig: a WorldConfig param[in] bufferManager: a BufferManager only used to allocate page blocks

void setDataTypeCoordinated(
LoraCache &other,
tensorrt_llm::DataType dataType,
SizeType32 totalNumPages,
SizeType32 otherTotalNumPages,
)#

Reinitialize this cache and another cache (e.g. a host/device pair) to store the given data type and page capacities as a single atomic operation.

Holds both caches’ mutexes for the full duration so a concurrent copyTask (which locks the same mutexes to verify the two caches agree on dtype before copying) can never observe one cache already reconfigured to the new dtype while the other still holds the old one.

tensorrt_llm::DataType getDataType() const#
void put(
TaskIdType taskId,
TensorPtr weights,
TensorPtr config,
bool load = true,
)#

put a task in the cache, and claim pages for it, and optionally load task weights.

Parameters:

  • taskId – [in] the task id

  • weights – [in] lora weights tensor

  • config – [in] lora config tensor

  • load – [in] if true load weights before returning, otherwise do not

void loadWeights(
TaskIdType taskId,
TensorPtr weights,
TensorPtr config,
)#

load task weights. This method must be called after put. It is designed to be called asynchronously after put returns with load = false

Parameters:

  • taslId – [in] the task id

  • weights – [in] lora weights tensor

  • config – [in] lora config tensor

inline bool isLoaded(TaskIdType taskId) const#
Parameters:

taskId – [in] the task id

Returns:

&#8212; true if task is loaded (weights are in place) and false otherwise

bool isDone(TaskIdType taskId) const#
Parameters:

taskId – [in] the task id

Returns:

&#8212; true if task is marked done and can be evicted

inline bool has(TaskIdType taskId) const#
Parameters:

taskId – [in] the task id

Returns:

&#8212; true if task is in the cache (not necessarily loaded) and false otherwise

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

taskId – [in] the task id

Returns:

&#8212; list of Value objects with pointers to task weights

void bump(TaskIdType taskId)#

bump task and make it the most recently used

Parameters:

taskId – [in] the task id

void markTaskDone(TaskIdType taskId)#

mark task done meaning it can be evicted

Parameters:

taskId – [in] the task id

void markAllDone()#

mark all tasks in cache done

SizeType32 determineNumPages(TaskIdType taskId) const#
Parameters:

taskId – [in] the taskid

Returns:

&#8212; number of pages needed to store the given task

SizeType32 determineNumPages(TensorPtr config) const#
Parameters:

config – [in] lora config tensor

Returns:

&#8212; number of pages needed to store the task configured with config tensor

bool fits(TensorPtr config) const#
Parameters:

config – [in] a lora config tensor

Returns:

&#8212; true in task fits in cache false otherwise

void copyTask(
TaskIdType taskId,
LoraCache &deviceCache,
bool markDone = false,
)#

copy task to another cache. Caches must have the same page size.

Parameters:

  • taskId – [in] the task id to copy

  • otherCache – [in] the LoraCache to move the task to

  • markDone – [in] mark the copied task done as it’s copied

SizeType32 getNumPages() const#
Returns:

&#8212; total number of pages allocated to cache (used or not)

ITensor::SharedConstPtr getPagePtr(size_t pageId) const#
Parameters:

pageId – [in] the page id

Returns:

&#8212; const pointer to page

Public Static Functions

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,
)#

Copy task weights to cache pages.

Parameters:

  • weights – [in] task weights

  • config – [in] task config tensor

  • modelConfig – [in] a ModelConfig

  • worldConfig – [in] a WorldConfig

  • modelIdToModel – [in] map from lora module id to LoraModule

  • manager – [in] a BufferManager the manager to use to perform the copies

  • pages – [out] list of page tensors to copy weights to

  • pageIds – [in] page ids for the pages

Returns:

&#8212; list of cache Values objects

static void splitTransposeCpu(
ITensor &output,
ITensor const &input,
SizeType32 tpSize,
SizeType32 tpRank,
)#

splits second dim of input into tpSize parts and writes the tpRank split to output

Parameters:

  • output – [out] output tensor

  • input – [in] input tensor

  • tpSize – [in] number of splits

  • tpRank – [in] the split to write to output

Private Types

enum ValueStatus#

Values:

enumerator kVALUE_STATUS_MISSING#
enumerator kVALUE_STATUS_PROCESSING#
enumerator kVALUE_STATUS_LOADED#
using TaskValuePtr = std::shared_ptr<TaskValue>#

Private Functions

void loadWeights(
TaskValue &cacheValue,
TensorPtr weights,
TensorPtr config,
)#
void bumpTaskInProgress(TaskIdType taskId)#
ValueStatus getStatus(TaskIdType taskId) const#
void setDataTypeLocked(
tensorrt_llm::DataType dataType,
std::optional<SizeType32> totalNumPages = std::nullopt,
)#

Core of setDataType, assumes mPagesMutex and mCacheMutex are already held by the caller. Used by setDataTypeCoordinated to reconfigure two caches under one combined lock.

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

claim numPages, evicting tasks if needed

Parameters:

numPages – [in] number of pages to claim

Throws:

std::runtime_error – if all pages cannot be claimed

Returns:

&#8212; list of page ids

std::map<size_t, std::pair<size_t, SizeType32>> copyTaskMapPages(
TaskValue &targetTaskValue,
TaskValue const &sourceTaskValue,
std::vector<size_t> const &targetPageIds,
LoraCache const &targetCache,
)#

Internal helper method used inside copyTask. Not thread safe on its own

Private Members

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#

Private Static Functions

template<typename T>
static void splitTransposeCpuInner(
ITensor &output,
ITensor const &input,
SizeType32 tpSize,
SizeType32 tpRank,
)#
struct TaskLayerModuleConfig#
#include <loraCache.h>

Contains information on a single layer / module. A list of these configs is associated with each task and can be used to populate runtime tensors.

Public Functions

std::string toString() const#
bool operator==(LoraCache::TaskLayerModuleConfig const &o) const#

Public Members

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#

Friends

friend class TaskLayerModuleConfigBindings
struct TaskValue#

Holds configuration and state for a single task.

Public Functions

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)#

Public Members

std::vector<std::size_t> pageIds#
TaskLayerModuleConfigListPtr configs#
std::list<TaskIdType>::iterator it#
bool inProgress#
bool loaded#
bool done#

Marks a task a done. This is used to mark a task as done during loading. if done=true at the end of loading (end of put, loadweights, or copyTask) the task will be marked as done

bool loadInProgress#

Indicates weights are loading either in put or loadWeights This is used to block concurrent loadWeights calls for the same task.

class LoraCacheFullException : public tensorrt_llm::runtime::LoraExpectedException#

Public Functions

explicit LoraCacheFullException(std::string const &msg)#
~LoraCacheFullException() noexcept override#
class LoraCachePageManager#
#include <loraCache.h>

Holds memory of lora cache pages, and manages allocation and freeing of whole pages. Memory is pre-allocated either on the host or device

Note that this class is not thread safe

Public Types

using TensorPtr = ITensor::SharedPtr#

Public Functions

LoraCachePageManager(
LoraCachePageManagerConfig const &config,
BufferManager const &bufferManager,
)#
Parameters:
std::optional<std::vector<std::size_t>> claimPages(
SizeType32 numPages,
)#

claim pages

Parameters:

numPages – [in] number of pages to claim

Returns:

a tuple, where the first values is a boolean indicating whether pages were claimed. If the first value is true the second value will have a list of pageIds

SizeType32 numAvailablePages() const#

get number of available (free) pages in manager

Returns:

number of free pages in manager

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

release given pages

Parameters:

pages – [in] list of pages to release (free)

ITensor::SharedConstPtr blockPtr(SizeType32 blockIdx) const#

return pointer to given page block

Parameters:

blockIdx; – [in]

Returns:

&#8212; pointer to page block

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

return pointer to given page

Parameters:

pageIdx – [in]

Returns:

&#8212; const pointer to page

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

return pointer to given page

Parameters:

pageIdx – [in]

Returns:

&#8212; mutable pointer to page

Private Functions

void initialize(BufferManager const &bufferManager)#

Private Members

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#

Subclassed by tensorrt_llm::runtime::LoraCacheFullException

Public Functions

explicit LoraExpectedException(std::string const &msg)#
~LoraExpectedException() noexcept override#

loraCachePageManagerConfig.h#

namespace tensorrt_llm
namespace runtime

Functions

inline std::ostream &operator<<(
std::ostream &os,
LoraCachePageManagerConfig const &c,
)#
inline std::string to_string(LoraCachePageManagerConfig const &c)#
class LoraCachePageManagerConfig#
#include <loraCachePageManagerConfig.h>

Configuration for LoraCachePageManager

See LoraCache docs for description of pages, slots, and page blocks.

Public Functions

inline explicit constexpr LoraCachePageManagerConfig(
runtime::MemoryType memType,
tensorrt_llm::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 tensorrt_llm::DataType constexpr getDataType() const noexcept#
inline void constexpr setDataType(
tensorrt_llm::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#

Private Members

runtime::MemoryType mMemoryType#
tensorrt_llm::DataType mDataType#
SizeType32 mTotalNumPages#
SizeType32 mMaxPagesPerBlock#
SizeType32 mSlotsPerPage#
SizeType32 mPageWidth#
SizeType32 mNumCopyStreams = 1#
bool mInitToZero#

speculativeDecodingModule.h#

namespace tensorrt_llm
namespace runtime
class SpeculativeDecodingModule#

Public Functions

inline explicit SpeculativeDecodingModule(
SizeType32 maxDraftPathLen,
SizeType32 maxDecodingDraftTokens,
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#
Returns:

max number of draft tokens that can be accepted by one step of the decoder

inline SizeType32 getMaxPathLen() const noexcept#

one more than draft path len for prediction from primary head

Returns:

max number of tokens that a request can grow in one step of the decoder

inline SizeType32 getMaxDecodingDraftTokens() const noexcept#
Returns:

max number of draft tokens processed by one step of the decoder

inline SizeType32 getMaxDecodingTokens() const noexcept#

one more than decoding draft tokens for prediction from primary head

Returns:

max number of tokens processed by one step of the decoder

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#

Private Functions

inline void computeNumPackedMasks() noexcept#

Private Members

SizeType32 mMaxDraftPathLen#
SizeType32 mMaxDecodingDraftTokens#
SizeType32 mMaxNumPaths#
SizeType32 mMaxNumPackedMasks#

iBuffer.h#

namespace tensorrt_llm
namespace runtime

Typedefs

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

Enums

enum class MemoryType : std::int32_t#

Values:

enumerator kGPU#
enumerator kCPU#
enumerator kPINNED#
enumerator kUVM#
enumerator kPINNEDPOOL#

Functions

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)#

Gets a typed pointer to the constant underlying data of the buffer.

Template Parameters:

T – The type of the underlying data.

Parameters:

buffer – The buffer to get a pointer to.

Returns:

A pointer to constant T.

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

Gets a typed pointer to the underlying data of the buffer.

Template Parameters:

T – The type of the underlying data.

Parameters:

buffer – The buffer to get a pointer to.

Returns:

A pointer to T.

template<typename T>
T *bufferCastOrNull(
IBuffer::SharedPtr const &bufferPtr,
)#

Retrieves a T typed pointer to the underlying data of the buffer pointed to by the bufferPtr, or nullptr if the bufferPtr is null.

Template Parameters:

T – The type of the underlying data.

Parameters:

bufferPtr – A possibly null shared ptr.

Returns:

A pointer to T, possibly nullptr.

template<typename T>
T const *bufferCastOrNull(
IBuffer::SharedConstPtr const &bufferPtr,
)#

Retrieves a T const typed pointer to the underlying data of the buffer pointed to by the bufferPtr, or nullptr if the bufferPtr is null.

Template Parameters:

T – The type of the underlying data.

Parameters:

bufferPtr – A possibly null shared ptr.

Returns:

A pointer to const T, possibly nullptr.

template<typename T>
T *bufferCastOrNull(
std::optional<IBuffer::SharedPtr> const &optionalBufferPtr,
)#

Retrieves a T typed pointer to the underlying data of the buffer pointed to by the buffer pointer contained in the optionalBufferPtr, or nullptr if the optional doesn’t have a value.

Template Parameters:

T – The type of the underlying data.

Parameters:

optionalBufferPtr – A possibly empty optional.

Returns:

A pointer to T, possibly nullptr.

template<typename T>
T const *bufferCastOrNull(
std::optional<IBuffer::SharedConstPtr> const &optionalBufferPtr,
)#

Retrieves a T const typed pointer to the underlying data of the buffer pointed to by the buffer pointer contained in the optionalBufferPtr, or nullptr if the optional doesn’t have a value.

Template Parameters:

T – The type of the underlying data.

Parameters:

optionalBufferPtr – A possibly empty optional.

Returns:

A pointer to const T, possibly nullptr.

std::ostream &operator<<(
std::ostream &output,
IBuffer const &buffer,
)#

Utility function to print a buffer.

class BufferDataType#
#include <iBuffer.h>

A wrapper around tensorrt_llm::DataType that provides a support for pointer types.

Public Functions

inline constexpr BufferDataType(
tensorrt_llm::DataType dataType,
bool _unsigned = false,
bool pointer = false,
)#
inline constexpr operator tensorrt_llm::DataType() const noexcept#
inline constexpr tensorrt_llm::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#

Public Static Attributes

static auto constexpr kTrtPointerType = tensorrt_llm::DataType::kINT64#

Private Members

tensorrt_llm::DataType mDataType#
bool mUnsigned#
bool mPointer#
template<typename T>
class BufferRange : public tensorrt_llm::common::ArrayView<T>#

Public Types

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

Public Functions

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<tensorrt_llm::DataType kDataType, bool kIsUnsigned = false, bool kIsPointer = false>
struct DataTypeTraits#
#include <iBuffer.h>

For converting a TensorRT data type to a C++ data type.

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

Public Types

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

Public Static Attributes

static char constexpr name[] = "*"#
static auto constexpr size = sizeof(type)#
template<bool kUnsigned>
struct DataTypeTraits<tensorrt_llm::DataType::kBOOL, kUnsigned>#

Public Types

using type = bool#

Public Static Attributes

static char constexpr name[] = "bool"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kFLOAT>#

Public Types

using type = float#

Public Static Attributes

static char constexpr name[] = "float"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kHALF>#

Public Types

using type = half#

Public Static Attributes

static char constexpr name[] = "half"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kINT32>#

Public Types

using type = std::int32_t#

Public Static Attributes

static char constexpr name[] = "int32"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kINT32, true>#

Public Types

using type = std::uint32_t#

Public Static Attributes

static char constexpr name[] = "uint32"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kINT64>#

Public Types

using type = std::int64_t#

Public Static Attributes

static char constexpr name[] = "int64"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kINT64, true>#

Public Types

using type = std::uint64_t#

Public Static Attributes

static char constexpr name[] = "uint64"#
static auto constexpr size = sizeof(type)#
template<>
struct DataTypeTraits<tensorrt_llm::DataType::kINT8>#

Public Types

using type = std::int8_t#

Public Static Attributes

static char constexpr name[] = "int8"#
static auto constexpr size = sizeof(type)#
template<bool kUnsigned>
struct DataTypeTraits<tensorrt_llm::DataType::kUINT8, kUnsigned>#

Public Types

using type = std::uint8_t#

Public Static Attributes

static char constexpr name[] = "uint8"#
static auto constexpr size = sizeof(type)#
class IBuffer#

Subclassed by tensorrt_llm::runtime::ITensor

Public Types

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 = tensorrt_llm::DataType#

Public Functions

virtual void *data() = 0#

Returns a pointer to underlying array.

virtual void const *data() const = 0#

Returns a pointer to underlying array.

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

Returns a pointer to the underlying array at a given element index.

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

Returns a pointer to the underlying array at a given element index.

virtual std::size_t getSize() const = 0#

Returns the size (in number of elements) of the buffer.

inline virtual std::size_t getSizeInBytes() const#

Returns the size (in bytes) of the buffer.

virtual std::size_t getCapacity() const = 0#

Returns the capacity of the buffer.

virtual DataType getDataType() const = 0#

Returns the data type of the buffer.

virtual char const *getDataTypeName() const#
virtual MemoryType getMemoryType() const = 0#

Returns the memory type of the buffer.

virtual char const *getMemoryTypeName() const#
virtual void resize(std::size_t newSize) = 0#

Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity.

virtual void release() = 0#

Releases the buffer. It will be reset to nullptr.

virtual ~IBuffer() = default#
IBuffer(IBuffer const&) = delete#

Not allowed to copy.

IBuffer &operator=(IBuffer const&) = delete#

Not allowed to copy.

Public Static Functions

static char const *getDataTypeName(DataType dataType)#
static UniquePtr slice(
SharedPtr buffer,
std::size_t offset,
std::size_t size,
)#

Creates a sliced view on the underlying buffer. The view will have the same data type as buffer.

Parameters:

  • buffer – The buffer to view.

  • offset – The offset of the view.

  • size – The size of the view.

Returns:

A view on the 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)#

Returns a view on the underlying tensor which can be independently resized.

Parameters:

tensor – The tensor to view.

Returns:

A view on the tensor.

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

Returns a view on the underlying tensor with a different size.

Parameters:

  • tensor – The tensor to view.

  • size – The size of the view.

Returns:

A view on the 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,
)#

Wraps the given data in an IBuffer. The IBuffer will not own the underlying data and cannot be resized beyond capacity.

Parameters:

  • data – The data to wrap.

  • type – The data type of the data.

  • size – The size of the buffer.

  • capacity – The capacity of the buffer.

Returns:

An 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)#

Determine the memory type of a pointer.

Protected Functions

IBuffer() = default#
inline std::size_t toBytes(std::size_t size) const#

Returns an array index or size in bytes.

template<MemoryType T>
struct MemoryTypeString#
template<>
struct MemoryTypeString<MemoryType::kCPU>#

Public Static Attributes

static auto constexpr value = "CPU"#
template<>
struct MemoryTypeString<MemoryType::kGPU>#

Public Static Attributes

static auto constexpr value = "GPU"#
template<>
struct MemoryTypeString<MemoryType::kPINNED>#

Public Static Attributes

static auto constexpr value = "PINNED"#
template<>
struct MemoryTypeString<MemoryType::kPINNEDPOOL>#

Public Static Attributes

static auto constexpr value = "PINNEDPOOL"#
template<>
struct MemoryTypeString<MemoryType::kUVM>#

Public Static Attributes

static auto constexpr value = "UVM"#
template<typename T, bool = false>
struct TRTDataType#
#include <iBuffer.h>

For converting a C++ data type to a TensorRT data type.

template<>
struct TRTDataType<bool>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kBOOL#
template<>
struct TRTDataType<float>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kFLOAT#
template<>
struct TRTDataType<half>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kHALF#
template<>
struct TRTDataType<kernels::KVCacheIndex>#

Public Static Attributes

static constexpr auto value = TRTDataType<kernels::KVCacheIndex::UnderlyingType>::value#
template<>
struct TRTDataType<runtime::RequestType>#

Public Static Attributes

static constexpr auto value = TRTDataType<std::underlying_type_t<runtime::RequestType>>::value#
template<>
struct TRTDataType<std::int32_t>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kINT32#
template<>
struct TRTDataType<std::int64_t>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kINT64#
template<>
struct TRTDataType<std::int8_t>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kINT8#
template<>
struct TRTDataType<std::uint32_t>#

Public Static Attributes

static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT32, true}#
template<>
struct TRTDataType<std::uint64_t>#

Public Static Attributes

static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT64, true}#
template<>
struct TRTDataType<std::uint8_t>#

Public Static Attributes

static constexpr auto value = tensorrt_llm::DataType::kUINT8#
template<typename T>
struct TRTDataType<T*>#

Public Static Attributes

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

Private Static Attributes

static auto constexpr kUnderlyingType = BufferDataType{TRTDataType<std::remove_const_t<T>, false>::value}#
template<>
struct TRTDataType<void*>#

Public Static Attributes

static constexpr auto value = BufferDataType::kTrtPointerType#
On this page