Architecture
Public Functions
+Constructs a DisaggExecutorOrchestrator object.
+ctxEnginePaths – A vector of file paths to context engine files.
genEnginePaths – A vector of file paths to generation engine files.
ctxExecutorConfigs – A vector of ExecutorConfig for context executors.
genExecutorConfigs – A vector of ExecutorConfig for generation executors.
hasContextAwaitThreads – Whether or not there are threads that receive response for each generation executor.
hasGenAwaitThreads – Whether or not there are threads that receive response for each generation executor.
Enqueue context-only requests to context executors.
+requests – A vector of context-only requests.
selectContextId – The index of the context executor to use. If std::nullopt, the executor that has the smallest number of inflight requests will be used.
batch – If true,enqueue requests in same context executor.If false, will try to use a different executor for each request.
A vector of global request ids, corresponding to the order of the requests in requests, the id returned may be different from the request id in each executor.
Enqueue generation-only requests to generation executors.
+requests – A vector of generation-only requests.
globalRequestIds – A vector of global request ids, corresponding to the order of the requests,and must be the ids returned by the enqueueContext function.
selectGenIdx – The index of the generation executor to use. If std::nullopt, the executor that has the smallest number of inflight requests will be used.
batch – If true,enqueue requests in same generation executor.If false, will try to use a different executor for each request.
Await for context responses.
+timeout – The maximum time to wait for new responses
contextIdx – The index of the context executor to use. If std::nullopt, return ready responses in all context executors,if hasContextAwaitThreads is true, then this parameter must be std::nullopt.
A vector of responses with corresponding global request ids
+Await for generation responses.
+timeout – The maximum time to wait for new responses.
genIdx – The index of the generation executor to use. If std::nullopt, return ready responses in all generation executors,if hasGenAwaitThreads is true, then this parameter must be std::nullopt.
A vector of responses with corresponding global request ids.
+Indicates if the current process is allowed to enqueueRequests.
+Get context executors.
+Get generation executors.
+Private Members
+Public Functions
+Public Types
+Public Types
+Public Functions
+Returns a pointer to underlying array.
+Returns a pointer to underlying array.
+Returns the memory type of the buffer.
+Returns the number of elements in the tensor.
+Returns the size of the tensor in bytes.
+Set the entire memory to zero.
+stream – Must be a valid CUDA stream if the memory type is GPU.
+Copy the data and shape from another tensor.
+other – A tensor to copy from.
stream – Must be a valid CUDA stream if the memory type is GPU.
Public Static Functions
+Allocate a cpu tensor with the given shape and data type.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a cpu tensor in pinned memory with the given shape and data type.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a cpu tensor in pooled pinned memory with the given shape and data type.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a tensor in managed memory (UVM) with the given shape and data type.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a gpu tensor with the given shape and data type on a particular cuda stream.
+shape – The shape of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
dataType – The data type of the tensor.
Wrap a data pointer into a tensor without taking ownership.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Wrap a data pointer into a tensor without taking ownership.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Wrap any container into a tensor without taking ownership.
+shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Private Functions
+ + +Public Static Functions
+Public Functions
-Public Functions
-Public Functions
-Constructs a DisaggExecutorOrchestrator object.
-ctxEnginePaths – A vector of file paths to context engine files.
genEnginePaths – A vector of file paths to generation engine files.
ctxExecutorConfigs – A vector of ExecutorConfig for context executors.
genExecutorConfigs – A vector of ExecutorConfig for generation executors.
hasContextAwaitThreads – Whether or not there are threads that receive response for each generation executor.
hasGenAwaitThreads – Whether or not there are threads that receive response for each generation executor.
Enqueue context-only requests to context executors.
-requests – A vector of context-only requests.
selectContextId – The index of the context executor to use. If std::nullopt, the executor that has the smallest number of inflight requests will be used.
batch – If true,enqueue requests in same context executor.If false, will try to use a different executor for each request.
A vector of global request ids, corresponding to the order of the requests in requests, the id returned may be different from the request id in each executor.
Enqueue generation-only requests to generation executors.
-requests – A vector of generation-only requests.
globalRequestIds – A vector of global request ids, corresponding to the order of the requests,and must be the ids returned by the enqueueContext function.
selectGenIdx – The index of the generation executor to use. If std::nullopt, the executor that has the smallest number of inflight requests will be used.
batch – If true,enqueue requests in same generation executor.If false, will try to use a different executor for each request.
Await for context responses.
-timeout – The maximum time to wait for new responses
contextIdx – The index of the context executor to use. If std::nullopt, return ready responses in all context executors,if hasContextAwaitThreads is true, then this parameter must be std::nullopt.
A vector of responses with corresponding global request ids
-Await for generation responses.
-timeout – The maximum time to wait for new responses.
genIdx – The index of the generation executor to use. If std::nullopt, return ready responses in all generation executors,if hasGenAwaitThreads is true, then this parameter must be std::nullopt.
A vector of responses with corresponding global request ids.
-Indicates if the current process is allowed to enqueueRequests.
-Get context executors.
-Get generation executors.
-Private Members
-Public Functions
-Public Types
-Public Types
-Public Functions
-Returns a pointer to underlying array.
-Returns a pointer to underlying array.
-Returns the memory type of the buffer.
-Returns the number of elements in the tensor.
-Returns the size of the tensor in bytes.
-Set the entire memory to zero.
-stream – Must be a valid CUDA stream if the memory type is GPU.
-Copy the data and shape from another tensor.
-other – A tensor to copy from.
stream – Must be a valid CUDA stream if the memory type is GPU.
Public Static Functions
-Allocate a cpu tensor with the given shape and data type.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a cpu tensor in pinned memory with the given shape and data type.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a cpu tensor in pooled pinned memory with the given shape and data type.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a tensor in managed memory (UVM) with the given shape and data type.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
Allocate a gpu tensor with the given shape and data type on a particular cuda stream.
-shape – The shape of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
dataType – The data type of the tensor.
Wrap a data pointer into a tensor without taking ownership.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Wrap a data pointer into a tensor without taking ownership.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Wrap any container into a tensor without taking ownership.
-shape – The shape of the tensor.
dataType – The data type of the tensor.
stream – Specifies the CUDA stream on which to allocate the tensor for GPU memory.
Private Functions
- - -Public Static Functions
+Public Functions
Public Functions
+Public Types
-Public Functions
-Public Members
-Mandatory parameters.
-GPT decoder class with support for in-flight batching.
-Subclassed by tensorrt_llm::runtime::GptDecoderBatched
-Public Types
-Public Functions
-Setup the decoder before calling forward()
Disable Lookahead decoding.
-Run one step for all requests without blocking the host process and return the token for synchronization.
-Run one step for all requests and wait for completion on the host.
-Gather final beam search results for request batchIdx. Result will only be available after event returned.
Public Types
Public Functions
Public Members
[maxDecoderSteps][batchSize][1, beamWidth, vocabSizePadded], on gpu
-Mandatory parameters Logits
+Mandatory parameters Previously generated token ids for all steps before DecodingInput.step, [BS, BM, MSL]
Maximum number of decoding tokens of active slots.
+The tokens computed during the gatherTree step, [BS, BM, MSL] Necessary for “Streaming + Beam Search” mode since beam search kernels store ungathered tokens in ids.
Batch of active decoder slots, sorted by slots, [maxDecoderSteps][batchSize].
+New tokens at each generated token of maxTokensPerStep, [maxTokensPerStep, BS, BM].
A Vector of views on newTokensSteps for each token [BS, BM].
+Optional parameters FinishedState by decoding if any of the stop conditions are met or if DecodingInput.finished is true, [BS, BM]
+Mandatory parameters for Beam Search log-probility of generated tokens, [BS, BM, MSL], float
+Public Static Attributes
+Public Functions
+Public Functions
-Creates a new cuda stream on the current device. The stream will be destroyed in the destructor.
-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.
Pass an existing cuda stream to this object.
-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.
Construct with an existing cuda stream or the default stream by passing nullptr.
-Returns the device on which the stream was created.
-Returns the stream associated with this object.
-Synchronizes the stream.
-Private Types
+Public Types
Public Types
+ + +Public Functions
+A helper class for managing memory on host and device.
+Public Types
+ + + + +Public Functions
+Construct a BufferManager.
+cudaStream – [in] The cuda stream to use for all operations on GPU (allocation, de-allocation, copying, etc.).
+Destructor.
+Allocates an IBuffer of the given size on the GPU, using cudaMallocAsync.
Allocates an ITensor of the given dimensions on the GPU, using cudaMallocAsync.
Allocates an IBuffer of the given size and memory type.
Allocates an ITensor of the given dimensions and memory type.
Create an empty IBuffer of the given memory type. It may be resized later.
Create an empty ITensor of the given memory type. It may be reshaped later.
Set the contents of the given buffer to value.
Copy src to dst.
Copy src to dst.
Copy src into a new IBuffer with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Copy src into a new IBuffer with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Get the underlying cuda stream.
+The current size of the memory reserved by the memory pool.
+The current size of the memory used by the memory pool.
+The current size of the memory free in the memory pool.
+Try to trim the memory reserved by the pool to size bytes. This synchronizes implicitly with the stream.
Public Static Functions
+Allocates an IBuffer of the given size on the GPU, using cudaMalloc.
Allocates an ITensor of the given dimensions on the GPU, using cudaMalloc.
Allocates an IBuffer of the given size on the CPU.
Allocates an ITensor of the given dimensions on the CPU.
Allocates a pinned IBuffer of the given size on the CPU.
Allocates a pinned ITensor of the given dimensions on the CPU.
Allocates a pinned IBuffer of the given size on the CPU in the default memory pool.
Allocates a pinned ITensor of the given dimensions on the CPU in the default memory pool.
Allocates an IBuffer of the given size in UVM.
Allocates an ITensor of the given dimensions in UVM.
Allocates an ITensor of the given dimensions for NVLS.
Private Members
+Friends
+Public Functions
+Public Static Functions
+Private Members
+Public Functions
+Public Members
+Public Types
+ + +Public Functions
+Private Members
+ + +Functions
+Helper function to produce batch slots [0, 1, …, batchSize - 1] for paths that do not explicitly provide batch slots to the decoder.
+Public Types
+Public Functions
+explicitDraftTokensDType – is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder.
+Private Members
+Subclassed by tensorrt_llm::runtime::GptDecoder< T >
+Public Types
+ + +Public Functions
+explicitDraftTokensDType – is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder.
+Public Static Functions
+Public Types
+Public Functions
+Public Members
+ + +Private Functions
+Private Members
+Public Members
+ + + + + + + + + + + + + + +Public Functions
+Public Members
+ + + + + + + + +[maxBatchSize, maxDecodingTokens] or [numSequences, maxDecodingTokens]
+[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
+[maxBatchSize, maxNumPaths, maxPathLen] or [numSequences, maxNumPaths, maxPathLen]
+[maxBatchSize, maxNumPaths, maxPathLen] or [numSequences, maxNumPaths, maxPathLen]
+[maxBatchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] or [numGenSequences, maxDecodingTokens, ceil(maxDecodingTokens / 32)]
+[maxBatchSize * maxDecodingTokens] or [numSequences * maxDecodingTokens]
+[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
+[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
+[maxBatchSize, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] or [numSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens]
+Public Functions
+Functions
Public Members
-Public Types
-Values:
+Public Functions
Setup buffers for the decoder excluding speculative decoding.
-Private Types
+Private Members
+Public Functions
+Setup buffers for the cache indirection.
-This is used for beam search on pipeline parallel ranks without a decoder.
-Setup buffers for speculative decoding.
-Disable lookahead decoding.
-[batchSize], number of finished sequences per request, on gpu
-[batchSize, beamWidth], FinishedState value, on gpu
-[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without padding, on gpu. In case of beam search, contains the ungathered data.
-batchIdx – index of the batch
-[maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without padding for request batchIdx, on gpu. In case of beam search, contains the ungathered data.
[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains gathered token ids without padding, on gpu.
-batchIdx – index of the batch
-[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains gathered token ids without padding for request batchIdx, on gpu.
[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains parent ids collected during beam search without padding, on gpu
-[batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu
-[maxBeamWidth], cumulative log probabilities (per beam), on gpu
-[batchSize, maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu
-[maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu
-[batchSize, maxBeamWidth], sequence lengths, on gpu
-batchIdx – index of the batch
-[maxBeamWidth], sequence lengths for request batchIdx, on gpu
Get maxTokensPerStep tokens generated in the last forward pass.
-[maxTokensPerStep, batchSize, maxBeamWidth], tokens generated in last forward pass, on gpu
-[batchSize, maxDraftTokens], predicted draft tokens for next step, on gpu
-[batchSize], predicted draft tokens lengths for previous step, on gpu
-[batchSize], predicted draft tokens lengths for next step, on gpu
-[batchSize + 1], exclusive sum of accepted draft token lengths, on gpu
-[batchSize, maxAcceptedDraftTokensPerStep], accepted paths packed into continuous tensor, on gpu
-[maxTokensPerStep, batchSize, beamWidth], finished states of type FinishedState, on gpu
-Get the number of tokens for all requests in the batch.
-The number of tokens for all requests in the batch.
-Get the number of tokens for a specific request in the batch.
-batchIdx – The index of the request in the batch.
-The number of tokens for the specified request.
-Set the number of tokens for a specific request in the batch.
+)#CudaVirtualMemoryAllocator::Configuration
batchIdx – The index of the request in the batch.
numTokens – The number of tokens for the specified request.
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.
Get the speculative decoding mode.
-Get the explicit draft tokens buffers.
-Get the eagle buffers.
-Get the lookahead buffers.
-Workspace for beam search in streaming mode.
-Get the generation steps for all requests in the batch.
-The generation steps for all requests in the batch.
-Set the generation steps for all requests in the batch.
-generationSteps – The generation steps for all requests in the batch.
-Stateful inputs for the decoder. Allocated for maxBatchSize slots.
-Stateful outputs for the decoder. Allocated for maxBatchSize slots.
-Private Functions
-Private Members
-Stateful inputs for the decoder. Allocated for maxBatchSize slots.
-Stateful outputs for the decoder. Allocated for maxBatchSize slots.
-[maxTokensPerStep, batchSize, beamWidth] finished states of type FinishedState for each generated token of maxTokensPerStep, on gpu
-Workspace for beam search in streaming mode.
-[batchSize], the num tokens of each request.
-Public Functions
-Public Types
- - -Public Functions
-Public Members
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -Public Functions
-Private Members
-Public Functions
-Public Members
-Public Functions
-Public Members
-Mandatory parameters Previously generated token ids for all steps before DecodingInput.step, [BS, BM, MSL]
-The tokens computed during the gatherTree step, [BS, BM, MSL] Necessary for “Streaming + Beam Search” mode since beam search kernels store ungathered tokens in ids.
New tokens at each generated token of maxTokensPerStep, [maxTokensPerStep, BS, BM].
-A Vector of views on newTokensSteps for each token [BS, BM].
-Optional parameters FinishedState by decoding if any of the stop conditions are met or if DecodingInput.finished is true, [BS, BM]
-Mandatory parameters for Beam Search log-probility of generated tokens, [BS, BM, MSL], float
-Public Static Attributes
Private Functions
+Private Members
+Friends
+CUDAVirtualMemoryChunk is a handle to a piece of CUDA memory allocation, providing the ability to release and rematerialize the allocation.
+Public Types
+Values:
+Public Functions
+ + +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.
+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.
+Test if this CUDAVirtualMemoryChunk is managing a memory block.
+Private Functions
+Private Members
+Private Static Attributes
+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
+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
Add memory to be managed by this manager.
+The memory and internal state will remain valid if any exception is thrown.
+ +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.
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.
+ +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.
Remove the memory from the manager.
+handle – The handle provided to add.
The CUDAVirtualMemory object. If the handle is unknown, an empty CUDAVirtualMemory will be returned.
+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.
tag – the tag to select target memories.
+Number of objects selected.
+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.
tag – the tag to select target memories.
+Number of objects selected.
+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.
+The handle list.
+Private Types
+ + +Private Functions
+Private Members
+LocalCreator creates memory allocation locally through cuMemCreate.
+Public Functions
+Public Members
- - - - - - - - - - - -MemsetConfigurator fills the memory with given value.
+Public Functions
+MulticastConfigurator binds the allocation handle to the given multicast object and offset.
+ + +OffloadConfigurator offload the content of the allocation to the backup storage when teardown, and restore the content on the following setup.
+Public Functions
+UnicastConfigurator maps the allocation handle into the specified unicast address range.
+Public Functions
+Public Types
+Public Functions
+Public Members
+Private Functions
+Public Members
+ + + + + + + + + + + + + + + + + + + + + + + + +Subclassed by tensorrt_llm::runtime::ExplicitDraftTokensBuffers::EngineInputs
+Public Functions
+Public Members
+ + + + + + + + +[maxBatchSize, maxNumPaths, maxPathDraftLen] or [numGenSequences, maxNumPaths, maxPathDraftLen]
+[maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen]
+[maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen]
+[maxBatchSize, maxNumPaths, maxPathDraftLen, vocabSize] or [numGenSequences, maxNumPaths, maxPathDraftLen, vocabSize]
+Functions
+Utility function to print a shape.
+Utility function to print a tensor with its shape.
+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.
+T – The type of the underlying data.
+tensorPtr – A possibly null shared ptr.
+A pointer to T const, possibly nullptr.
+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.
+T – The type of the underlying data.
+tensorPtr – A possibly null shared ptr.
+A pointer to T, possibly nullptr.
+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.
+T – The type of the underlying data.
+optionalBufferPtr – A possibly empty optional.
+A pointer to T, possibly nullptr.
+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.
+T – The type of the underlying data.
+optionalBufferPtr – A possibly empty optional.
+A pointer to const T, possibly nullptr.
+Public Types
+ + + + + + + + + + + + +Public Functions
+Returns the tensor n-th dimension. If n is negative, returns the (nbDims - n)th dimension. TODO: replace with constexpr parameter when moving to C++20.
+Sets the tensor dimensions. The new size of the tensor will be volume(dims)
Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity.
+Removes the given unit dimensions from this tensor.
+Adds a unit dimension at the specified position.
+Public Static Functions
+Returns the volume of the dimensions. Returns -1 if d.nbDims < 0.
Returns the volume of the dimensions. Throws if d.nbDims < 0.
Returns the strides of each dimemsion in a Shape.
+Removes the given unit dimension from shape.
shape – The shape to squeeze.
dim – The dimension that should be removed (“squeezed”).
A new shape without the unit dimension.
+Add a unit dimension to shape at the specified position.
shape – The shape to unsqueeze.
dim – The dimension where unit dimension should be added.
A new shape with the added unit dimension.
+Creates a sliced view on the underlying tensor. The view will have the same data type as tensor.
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.
A view on the buffer.
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.
Whenever – offset overflows or the last dimension offset+size overflows.
+A view of shape [size, the rest dimensions] or [size] when
+return the rest slices at the last dimension when size omitted.
offsetDims – specifies all dimensions.
+Just the block at the point, with shape of [the rest dimensions] or [1] when
+Returns a view on the underlying buffer (or tensor) with the given shape.
tensor – The tensor to view.
shape – The shape of the view.
A view on the tensor.
Returns a view on the underlying tensor which can be independently reshaped.
tensor – The tensor to view.
+A view on the tensor.
Returns a flattened view on the underlying tensor which can be independently reshaped.
tensor – The tensor to flatten.
sliceN – Slice the first N elements after flattening. -1 means take the whole flattened tensor.
A flatten view on the tensor.
Wraps the given data in an ITensor. The ITensor will not own the underlying data and cannot be reshaped beyond capacity.
data – The data to wrap.
type – The data type of the data.
shape – The shape of the tensor.
capacity – The capacity of the buffer.
An ITensor.
A convenience function to create a tensor shape with the given dimensions.
+A convenience function for converting a tensor shape to a string.
A convenience function to compare shapes.
+A convenience function to compare shapes.
+Protected Functions
+Friends
+Defines
+Typedefs
+Enums
+ + +Public Functions
+Public Members
+Functions
+Configuration for LoraCachePageManager
+See LoraCache docs for description of pages, slots, and page blocks.
+Public Functions
+Private Members
+Public Functions
+Is my rank the last rank in its pipeline?
+Public Static Functions
+Public Static Attributes
+Private Members
+Functions
+Public Types
+Values:
+Public Functions
+Public Static Functions
+Private Members
+Public Types
+Public Functions
+Public Static Functions
+Private Functions
+Private Members
+Private Static Attributes
+Public Types
+Public Functions
+Creates a new cuda event. The event will be destroyed in the destructor.
+flags – Flags for event creation. By default, event timing is disabled.
+Pass an existing cuda event to this object.
+event – The event to pass to this object.
ownsEvent – Whether this object owns the event and destroys it in the destructor.
Synchronizes the event.
+Private Types
+ + +Subclassed by tensorrt_llm::runtime::EagleModule, tensorrt_llm::runtime::LookaheadModule, tensorrt_llm::runtime::MedusaModule
+Public Functions
Is my rank the last rank in its pipeline?
+max number of draft tokens that can be accepted by one step of the decoder
+Public Static Functions
-+one more than draft path len for prediction from primary head
+max number of tokens that a request can grow in one step of the decoder
+max number of draft tokens processed by one step of the decoder
++one more than decoding draft tokens for prediction from primary head
+max number of tokens processed by one step of the decoder
+Public Static Attributes
-Private Functions
+Private Members
GPT decoder class with support for in-flight batching.
+Subclassed by tensorrt_llm::runtime::GptDecoderBatched
+Public Types
+Public Functions
+Setup the decoder before calling forward()
Disable Lookahead decoding.
+Run one step for all requests without blocking the host process and return the token for synchronization.
+Run one step for all requests and wait for completion on the host.
+Gather final beam search results for request batchIdx. Result will only be available after event returned.
Public Types
+Public Functions
+Public Members
+[maxDecoderSteps][batchSize][1, beamWidth, vocabSizePadded], on gpu
+Mandatory parameters Logits
+Maximum number of decoding tokens of active slots.
+Public Functions
+Private Members
+Public Types
-Public Functions
Public Members
-Private Functions
-Public Members
- - - - - - - - - - - - - - - - - - - - - - - - -Subclassed by tensorrt_llm::runtime::ExplicitDraftTokensBuffers::EngineInputs
-Public Functions
-Public Members
- - - - - - - - -[maxBatchSize, maxNumPaths, maxPathDraftLen] or [numGenSequences, maxNumPaths, maxPathDraftLen]
-[maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen]
-[maxBatchSize, maxNumPaths, maxPathLen] or [numGenSequences, maxNumPaths, maxPathLen]
-[maxBatchSize, maxNumPaths, maxPathDraftLen, vocabSize] or [numGenSequences, maxNumPaths, maxPathDraftLen, vocabSize]
-A helper class for managing memory on host and device.
-Public Types
- - - - -Public Functions
-Construct a BufferManager.
+)#Creates a new cuda stream on the current device. The stream will be destroyed in the destructor.
cudaStream – [in] The cuda stream to use for all operations on GPU (allocation, de-allocation, copying, etc.).
+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.
Destructor.
-Allocates an IBuffer of the given size on the GPU, using cudaMallocAsync.
Allocates an ITensor of the given dimensions on the GPU, using cudaMallocAsync.
Allocates an IBuffer of the given size and memory type.
Allocates an ITensor of the given dimensions and memory type.
Create an empty IBuffer of the given memory type. It may be resized later.
Create an empty ITensor of the given memory type. It may be reshaped later.
Set the contents of the given buffer to value.
Copy src to dst.
Copy src to dst.
Copy src into a new IBuffer with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Copy src into a new IBuffer with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Copy src into a new ITensor with a potentially different memory type.
Get the underlying cuda stream.
-The current size of the memory reserved by the memory pool.
-The current size of the memory used by the memory pool.
-The current size of the memory free in the memory pool.
-Try to trim the memory reserved by the pool to size bytes. This synchronizes implicitly with the stream.
Public Static Functions
-Allocates an IBuffer of the given size on the GPU, using cudaMalloc.
Allocates an ITensor of the given dimensions on the GPU, using cudaMalloc.
Allocates an IBuffer of the given size on the CPU.
Allocates an ITensor of the given dimensions on the CPU.
Allocates a pinned IBuffer of the given size on the CPU.
Allocates a pinned ITensor of the given dimensions on the CPU.
Allocates a pinned IBuffer of the given size on the CPU in the default memory pool.
Allocates a pinned ITensor of the given dimensions on the CPU in the default memory pool.
Allocates an IBuffer of the given size in UVM.
Allocates an ITensor of the given dimensions in UVM.
Allocates an ITensor of the given dimensions for NVLS.
Private Members
-Friends
-Functions
-Public Types
-Values:
-Public Functions
-Public Static Functions
-Private Members
-Public Types
-Public Functions
-Public Members
- - -Private Functions
-Private Members
-Public Members
- - - - - - - - - - - - - - -Public Functions
-Public Members
- - - - - - - - -[maxBatchSize, maxDecodingTokens] or [numSequences, maxDecodingTokens]
-[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
-[maxBatchSize, maxNumPaths, maxPathLen] or [numSequences, maxNumPaths, maxPathLen]
-[maxBatchSize, maxNumPaths, maxPathLen] or [numSequences, maxNumPaths, maxPathLen]
-[maxBatchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] or [numGenSequences, maxDecodingTokens, ceil(maxDecodingTokens / 32)]
-[maxBatchSize * maxDecodingTokens] or [numSequences * maxDecodingTokens]
-[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
-[maxBatchSize, maxDecodingDraftTokens] or [numSequences, maxDecodingDraftTokens]
-[maxBatchSize, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens] or [numSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens]
-Public Types
-Public Functions
-Public Static Functions
-Private Functions
-Private Members
-Private Static Attributes
-Public Types
- - -Public Functions
- - -Public Types
- - -Public Functions
-Functions
-Helper function to produce batch slots [0, 1, …, batchSize - 1] for paths that do not explicitly provide batch slots to the decoder.
-Public Types
-Public Functions
-Pass an existing cuda stream to this object.
+explicitDraftTokensDType – is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder.
+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.
Construct with an existing cuda stream or the default stream by passing nullptr.
+Returns the device on which the stream was created.
+Returns the stream associated with this object.
+Synchronizes the stream.
+Private Types
+Private Members
Subclassed by tensorrt_llm::runtime::GptDecoder< T >
-Public Types
- - -Public Functions
-explicitDraftTokensDType – is only used by ExplicitDraftTokens model to WAR the lack of bf16 decoder.
-Public Static Functions
-Public Functions
Public Static Functions
-Private Members
Defines
+Public Types
- - -Public Functions
Public Members
Private Types
+Private Functions
+ + +Public Types
+Public Functions
+Public Members
+Mandatory parameters.
+Public Functions
+Public Members
+Public Types
+ + +Public Functions
+Setup buffers for the decoder excluding speculative decoding.
+Setup buffers for the cache indirection.
+This is used for beam search on pipeline parallel ranks without a decoder.
+Setup buffers for speculative decoding.
+Disable lookahead decoding.
+[batchSize], number of finished sequences per request, on gpu
+[batchSize, beamWidth], finished states of type FinishedState, on gpu
+[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without padding, on gpu. In case of beam search, contains the ungathered data.
+batchIdx – index of the batch
+[maxBeamWidth, maxInputLength + maxNewTokens], contains input token ids and generated token ids without padding for request batchIdx, on gpu. In case of beam search, contains the ungathered data.
[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains gathered token ids without padding, on gpu.
+batchIdx – index of the batch
+[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], only used for beam search. It contains gathered token ids without padding for request batchIdx, on gpu.
[batchSize, maxBeamWidth, maxInputLength + maxNewTokens], contains parent ids collected during beam search without padding, on gpu
+[batchSize, maxBeamWidth], cumulative log probabilities (per beam), on gpu
+[maxBeamWidth], cumulative log probabilities (per beam), on gpu
+[batchSize, maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu
+[maxBeamWidth, maxSequenceLength], log probabilities (per beam), on gpu
+[batchSize, maxBeamWidth], sequence lengths, on gpu
+batchIdx – index of the batch
+[maxBeamWidth], sequence lengths for request batchIdx, on gpu
Get maxTokensPerStep tokens generated in the last forward pass.
+[maxTokensPerStep, batchSize, maxBeamWidth], tokens generated in last forward pass, on gpu
+[batchSize, maxDraftTokens], predicted draft tokens for next step, on gpu
+[batchSize], predicted draft tokens lengths for previous step, on gpu
+[batchSize], predicted draft tokens lengths for next step, on gpu
+[batchSize + 1], exclusive sum of accepted draft token lengths, on gpu
+[batchSize, maxAcceptedDraftTokensPerStep], accepted paths packed into continuous tensor, on gpu
+Get the number of tokens for all requests in the batch.
+The number of tokens for all requests in the batch.
+Get the number of tokens for a specific request in the batch.
+batchIdx – The index of the request in the batch.
+The number of tokens for the specified request.
+Set the number of tokens for a specific request in the batch.
+batchIdx – The index of the request in the batch.
numTokens – The number of tokens for the specified request.
Get the speculative decoding mode.
+Get the explicit draft tokens buffers.
+Get the eagle buffers.
+Get the lookahead buffers.
+Workspace for beam search in streaming mode.
+Get the generation steps for all requests in the batch.
+The generation steps for all requests in the batch.
+Set the generation steps for all requests in the batch.
+generationSteps – The generation steps for all requests in the batch.
+Stateful inputs for the decoder. Allocated for maxBatchSize slots.
+Stateful outputs for the decoder. Allocated for maxBatchSize slots.
+Private Functions
+Private Members
Stateful inputs for the decoder. Allocated for maxBatchSize slots.
+Stateful outputs for the decoder. Allocated for maxBatchSize slots.
+Workspace for beam search in streaming mode.
+[batchSize], the num tokens of each request.
+Typedefs
- - -Enums
- - -Functions
-Gets a typed pointer to the constant underlying data of the buffer.
-T – The type of the underlying data.
-buffer – The buffer to get a pointer to.
-A pointer to constant T.
Gets a typed pointer to the underlying data of the buffer.
-T – The type of the underlying data.
-buffer – The buffer to get a pointer to.
-A pointer to T.
Retrieves a T typed pointer to the underlying data of the buffer pointed to by the bufferPtr, or nullptr if the bufferPtr is null.
-T – The type of the underlying data.
-bufferPtr – A possibly null shared ptr.
-A pointer to T, possibly nullptr.
-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.
-T – The type of the underlying data.
-bufferPtr – A possibly null shared ptr.
-A pointer to const T, possibly nullptr.
-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.
-T – The type of the underlying data.
-optionalBufferPtr – A possibly empty optional.
-A pointer to T, possibly nullptr.
-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.
-T – The type of the underlying data.
-optionalBufferPtr – A possibly empty optional.
-A pointer to const T, possibly nullptr.
-A wrapper around nvinfer1::DataType that provides a support for pointer types.
Public Functions
-Public Static Attributes
- - -Public Types
-For converting a TensorRT data type to a C++ data type.
-Public Types
-Public Types
-Public Types
-Public Types
-Public Types
-Public Types
-Public Types
-Public Types
-Public Types
-Subclassed by tensorrt_llm::runtime::ITensor
-Public Types
- - - - - - - - - - -Public Functions
-Returns a pointer to underlying array.
-Returns a pointer to underlying array.
-Returns a pointer to the underlying array at a given element index.
-Returns a pointer to the underlying array at a given element index.
-Returns the size (in number of elements) of the buffer.
-Returns the size (in bytes) of the buffer.
-Returns the capacity of the buffer.
-Returns the memory type of the buffer.
-Resizes the buffer. This is a no-op if the new size is smaller than or equal to the current capacity.
-Releases the buffer. It will be reset to nullptr.
-Public Static Functions
- - -Creates a sliced view on the underlying buffer. The view will have the same data type as buffer.
buffer – The buffer to view.
offset – The offset of the view.
size – The size of the view.
A view on the buffer.
Returns a view on the underlying tensor which can be independently resized.
tensor – The tensor to view.
-A view on the tensor.
Returns a view on the underlying tensor with a different size.
tensor – The tensor to view.
size – The size of the view.
A view on the tensor.
Wraps the given data in an IBuffer. The IBuffer will not own the underlying data and cannot be resized beyond capacity.
data – The data to wrap.
type – The data type of the data.
size – The size of the buffer.
capacity – The capacity of the buffer.
An IBuffer.
Determine the memory type of a pointer.
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-For converting a C++ data type to a TensorRT data type.
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Public Static Attributes
-Private Static Attributes
-Public Static Attributes
-Public Functions
Public Static Functions
Private Members
Functions
-Configuration for LoraCachePageManager
-See LoraCache docs for description of pages, slots, and page blocks.
-Public Functions
-Private Members
-