Decoder Utils#

std::unique_ptr<EngineExecutor> trt_edgellm::rt::decoder_utils::loadDraftEngine(
std::filesystem::path const &engineDir,
DeploymentConfig &deployment,
)#

Load the draft engine from disk and return an EngineExecutor.

void trt_edgellm::rt::decoder_utils::appendAcceptedTokens(
DecodingInferenceContext &context,
Tensor &hostAcceptLengths,
Tensor &hostAcceptedTokenIds,
Tensor const &deviceAcceptLength,
Tensor const &deviceAcceptedTokenIds,
int32_t maxAcceptDepth,
tokenizer::Tokenizer const &tokenizer,
cudaStream_t stream,
)#

Copy accepted tokens from device buffers into the host-side context token lists. On return, hostAcceptLengths holds the number of tokens actually appended per slot.

void trt_edgellm::rt::decoder_utils::enqueueLogprobsD2H(
Tensor const &inputLogits,
int32_t rows,
DecodingRuntimeContext &runtime,
int32_t topK,
cudaStream_t stream,
)#

Enqueue log-softmax + top-K extraction + async D2H staging for one decode step. Device work only — results land in runtime.logprobs.host* after the caller’s round sync.

Parameters:
  • inputLogits – Row-major logits (GPU): [rows, vocabSize].

  • rows – Total rows: activeBatchSize (vanilla / prefill) or activeBatchSize * rowsPerBatch (spec decode, gathered if needed).

  • runtime – Runtime context providing logprobs/sampling buffers.

  • topK – Number of top log-probabilities to extract.

  • stream – CUDA stream (not synchronized here).

void trt_edgellm::rt::decoder_utils::collectLogprobsFromHost(
DecodingRuntimeContext &runtime,
DecodingInferenceContext &context,
int32_t activeBatchSize,
int32_t topK,
)#

Collect staged logprobs into context.stepLogprobs (vanilla / prefill: one row per slot). Call after the round synchronization that followed enqueueLogprobsD2H().

void trt_edgellm::rt::decoder_utils::collectSpecLogprobsFromHost(
DecodingRuntimeContext &runtime,
DecodingInferenceContext &context,
int32_t activeBatchSize,
int32_t rowsPerBatch,
int32_t const *hostAcceptLens,
int32_t topK,
)#

Collect staged logprobs into context.stepLogprobs (spec decode: acceptLen rows per slot). Call after the round synchronization (appendAcceptedTokens) that made hostAcceptLens valid.

Parameters:

rowsPerBatch – Max rows per batch item: maxAcceptDepth for EAGLE/MTP, blockSize for DFlash.