Decoder Utils

void trt_edgellm::rt::decoder_utils::zeroActiveRegion(

Tensor &tensor,

cudaStream_t stream

)

Zero the region tensor's current shape covers, not its whole allocation.

The spec-decode hidden-state buffers are sized for max_input_len but reshaped down to the few tokens a step binds, so clearing getMemoryCapacity() would cost O(max_input_len) per token. The bound region is a contiguous prefix of the allocation, so shape.volume() is sufficient.

void trt_edgellm::rt::decoder_utils::prepareSpecRaggedBindings(

DecodingRuntimeContext &runtime,

LLMEngineConfig const &cfg,

int32_t kvCacheIndex,

Tensor const &attentionPositions,

Tensor const &committedPastLengths,

Tensor const *validCounts,

Tensor const &selectedTokenIndices,

int32_t selectedRows,

std::vector<ResidentRef> const *residentRefs,

int32_t batchSize,

int32_t queryWidth,

InferenceDims const &dims,

cudaStream_t stream

)

RaggedExecutionBatch trt_edgellm::rt::decoder_utils::buildSpecPrefillRaggedBatch(

DecodingInferenceContext const &context,

int32_t queryWidth

)

ExecutionPhase trt_edgellm::rt::decoder_utils::contextPrefillPhase(

std::vector<int32_t> const &pastLengths,

int32_t activeBatchSize

)

void trt_edgellm::rt::decoder_utils::prepareSpecPrefillRaggedBindings(

DecodingRuntimeContext &runtime,

LLMEngineConfig const &cfg,

int32_t kvCacheIndex,

DecodingInferenceContext const &context,

int32_t queryWidth,

InferenceDims const &dims,

cudaStream_t stream

)

std::unique_ptr<EngineExecutor> trt_edgellm::rt::decoder_utils::loadDraftEngine(

std::filesystem::path const &engineDir,

DeploymentConfig const &deployment

)

Load the draft engine from disk and return an EngineExecutor.

void trt_edgellm::rt::decoder_utils::directVocabMapToOffsets(

std::vector<int32_t> &vocabMap,

int32_t baseVocabSize

)

Validate a DIRECT reduced-to-full vocabulary map and convert it in place to OFFSETS.

void trt_edgellm::rt::decoder_utils::appendSampledTokens(

DecodingInferenceContext &context,

int32_t const *sampledTokenIds,

int32_t activeBatchSize

)

Append one sampled token per active slot to the host-side token lists.

Slots already marked finished are skipped: they are evicted at the end of this step, so a token appended to them is never fed back and only pollutes the output. This mirrors collectLogprobsFromHost, which skips the same slots, keeping tokenIds and stepLogprobs in step with each other.

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,

int32_t proposedDraftsPerRound = 0,

int32_t const *perSlotProposedDrafts = nullptr

)

Copy accepted tokens from device buffers into the host-side context token lists. The caller must commit the same accept lengths to persistent base state before calling this function. On return, hostAcceptLengths holds the number of tokens actually appended per slot.

void trt_edgellm::rt::decoder_utils::clampAcceptLengthsToRemainingGeneration(

DecodingInferenceContext const &context,

Tensor &deviceAcceptLength,

cudaStream_t stream

)

Clamp device accept lengths so multi-token speculative commits never exceed max_generate_length.

Runs entirely on stream: the per-slot budgets travel as kernel arguments, so the verify round has no D2H copy and no host synchronization before appendAcceptedTokens(). A mid-round host sync drains the stream, and under GPU time-slicing that forfeits a full foreign quantum every round.

void trt_edgellm::rt::decoder_utils::applyForcedAcceptance(

DecodingInferenceContext &context,

Tensor &hostAcceptLengths,

Tensor &hostAcceptedTokenIds,

Tensor &deviceAcceptLength,

Tensor &deviceAcceptedTokenIds,

std::vector<int32_t> &ownTokens,

int32_t maxAcceptDepth,

cudaStream_t stream

)

Teacher-force this round’s acceptance to the golden’s tokens.

Call after clampAcceptLengthsToRemainingGeneration() and before the KV-cache commit, since trimming the acceptance is what keeps a replaced token’s stale cache entry out of the commit.

Parameters:

ownTokens – Out: per sequence, the token it would itself have committed at the slot that ends up last &#8212; the divergence signal the dump records.

void trt_edgellm::rt::decoder_utils::dumpSpecRound(

DecodingInferenceContext &context,

HybridCacheManager &cacheManager,

KVPageTable const &pageTable,

KVPageTable const *swaPageTable,

Tensor const &verifyLogits,

Tensor const &acceptedTokenIndices,

Tensor const &hostAcceptLengths,

std::vector<int32_t> const &ownTokens,

int32_t verifySize,

int32_t maxAcceptDepth,

cudaStream_t stream

)

Dump one speculative round’s committed base state.

Call after appendAcceptedTokens(), so the token list and the cache are both final.

Parameters:

• verifyLogits – Base verify logits [activeBatch, verifySize, vocab].

• acceptedTokenIndices – Device [activeBatch, maxAcceptDepth] verify rows that were accepted.

• hostAcceptLengths – Host accept lengths, as written back by appendAcceptedTokens().

• ownTokens – From applyForcedAcceptance(); empty when forcing is off.

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 (multi-row 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, or canvasLen for DiffusionGemma.