Sampling#

struct SamplingParams#

Structure to hold sampling parameters.

Public Functions

inline SamplingParams(
int32_t batchSize_,
int32_t vocabSize_,
float temperature_ = 1.0f,
int32_t topK_ = 0,
float topP_ = 1.0f,
)#

Constructor with default values.

Parameters:

  • batchSize_ – Number of samples in the batch

  • vocabSize_ – Size of the vocabulary

  • temperature_ – Temperature parameter (default: 1.0f)

  • topK_ – Top-K parameter (default: 0, disabled)

  • topP_ – Top-P parameter (default: 1.0f, disabled)

Throws:

std::invalid_argument – if neither topK nor topP is set, or if temperature is invalid

Public Members

int32_t batchSize#

Number of samples in the batch.

int32_t vocabSize#

Size of the vocabulary.

float temperature#

Temperature parameter for sampling (higher = more random)

int32_t topK#

Top-K sampling parameter (0 = disabled)

float topP#

Top-P (nucleus) sampling parameter (1.0 = disabled)

bool useTopK#

Flag indicating if top-K sampling is enabled.

bool useTopP#

Flag indicating if top-P sampling is enabled.

size_t trt_edgellm::getTopKtopPSamplingWorkspaceSize(
int32_t batchSize,
int32_t vocabSize,
SamplingParams const &params,
)#

Get workspace size required for top-K/top-P sampling (FP32 only).

Calculates the amount of GPU memory needed for intermediate computations during the sampling operation. The workspace must be allocated before calling topKtopPSamplingFromLogits().

Parameters:

  • batchSize – [in] Batch size for sampling

  • vocabSize – [in] Vocabulary size

  • params – [in] Sampling parameters

Throws:

std::runtime_error – if topK and topP are both not set

Returns:

Required workspace size in bytes

size_t trt_edgellm::getSelectAllTopKWorkspaceSize(
int32_t batchSize,
int32_t vocabSize,
int32_t topK,
) noexcept#

Get workspace size required for selectAllTopK operation (FP32 only).

Calculates the amount of GPU memory needed for intermediate computations during the top-K selection operation. The workspace must be allocated before calling selectAllTopK().

Parameters:

  • batchSize – [in] Batch size for selection

  • vocabSize – [in] Vocabulary size

  • topK – [in] Number of top elements to select

Returns:

Required workspace size in bytes

void trt_edgellm::applyLogitBias(
rt::Tensor &logits,
rt::Tensor const &tokenIds,
rt::Tensor const &biasValues,
rt::Tensor const &offsets,
cudaStream_t stream,
)#

Apply sparse per-batch logit biases in place before sampling.

Biases are represented in CSR-style flattened buffers. For each batch row b, entries in [offsets[b], offsets[b + 1]) are added to logits[b, tokenIds[i]]. Invalid token IDs are ignored defensively; callers should still validate token IDs before invoking this helper.

Parameters:

  • logits – [inout] Logits tensor [GPU, Float] with shape [batch-size, vocab-size]

  • tokenIds – [in] Flattened biased token IDs [GPU, Int32] with shape [num-biased-tokens]

  • biasValues – [in] Flattened bias values [GPU, Float] with shape [num-biased-tokens]

  • offsets – [in] Per-batch CSR offsets [GPU, Int32] with shape [batch-size + 1]

  • stream – [in] CUDA stream to execute the kernel

Throws:

std::runtime_error – If tensor validation or CUDA launch fails

void trt_edgellm::topKtopPSamplingFromLogits(
rt::Tensor const &logits,
rt::Tensor &selectedIndices,
SamplingParams const &params,
rt::Tensor &workspace,
cudaStream_t stream,
uint64_t philoxSeed = 42,
uint64_t philoxOffset = 0,
)#

Main sampling function for top-K and top-P sampling from logits.

Performs token sampling using top-K and/or top-P (nucleus) sampling strategies on the input logits. The function applies temperature scaling and returns the selected token indices for each batch element.

Parameters:

  • logits – [in] Input logits tensor [GPU, Float] with shape [batch-size, vocab-size]

  • selectedIndices – [out] Selected token indices [GPU, Int32] with shape [batch-size, 1]

  • params – [in] Sampling parameters including batch size, vocab size, temperature, top-K, and top-P values

  • workspace – [inout] Workspace buffer [GPU, Int8] for intermediate computations

  • stream – [in] CUDA stream to execute the kernel

  • philoxSeed – [in] Random seed for sampling (default: 42)

  • philoxOffset – [in] Random offset for sampling (default: 0)

Throws:

std::runtime_error – If CUDA operations fail

void trt_edgellm::selectAllTopK(
rt::Tensor const &input,
rt::OptionalOutputTensor topKValues,
rt::Tensor &topKIndices,
int32_t topK,
rt::Tensor &workspace,
cudaStream_t stream,
)#

Select all top-K elements from input tensor.

Returns topK indices and raw values from input with no transformations applied. This function identifies the K largest elements in each batch and returns their indices and optionally their values.

Parameters:

  • input – [in] Input tensor [GPU, Float] with shape [batch-size, vocab-size]

  • topKValues – [out] Optional top-K values [GPU, Float] with shape [batch-size, top-K]. Can be std::nullopt if values not needed

  • topKIndices – [out] Top-K indices [GPU, Int32] with shape [batch-size, top-K]

  • topK – [in] Number of top elements to select

  • workspace – [inout] Workspace buffer [GPU, Int8] for intermediate computations

  • stream – [in] CUDA stream to execute the kernel

Throws:

std::runtime_error – If CUDA operations fail

size_t trt_edgellm::getExtractTopKLogprobsWorkspaceSize(
int32_t batchSize,
int32_t vocabSize,
int32_t topK,
) noexcept#

Get workspace size required for extractTopKLogprobs.

The workspace holds the intermediate log-softmax buffer ([batch-size, vocab-size] floats, 256-byte aligned) followed by the selectAllTopK temporary storage. It is NOT a plain selectAllTopK workspace; always size logprobs workspaces with this function.

Parameters:

  • batchSize – [in] Batch size

  • vocabSize – [in] Vocabulary size

  • topK – [in] Number of top log-probabilities to extract

Returns:

Required workspace size in bytes

void trt_edgellm::extractTopKLogprobs(
rt::Tensor const &logits,
rt::Tensor &logprobsValues,
rt::Tensor &logprobsIndices,
int32_t topK,
rt::Tensor &workspace,
cudaStream_t stream,
)#

Extract top-K log-probabilities from raw logits.

Computes numerically-stable log-softmax over the full vocabulary for each batch element, then selects the top-K (token_id, log_prob) pairs sorted by descending probability. No temperature scaling is applied — logprobs match OpenAI convention: log(softmax(logits)).

Parameters:

  • logits – [in] Input logits [GPU, Float] with shape [batch-size, vocab-size]

  • logprobsValues – [out] Top-K log-prob values [GPU, Float] with shape [batch-size, top-K]

  • logprobsIndices – [out] Top-K token indices [GPU, Int32] with shape [batch-size, top-K]

  • topK – [in] Number of top elements to select (must be <= kMaxLogprobsK)

  • workspace – [inout] Workspace buffer [GPU, Int8]; holds the log-softmax buffer and the top-K selection temp storage. Size it with getExtractTopKLogprobsWorkspaceSize().

  • stream – [in] CUDA stream

Throws:

std::runtime_error – If CUDA operations fail

void trt_edgellm::gatherSpecVerifyAcceptedLogitRows(
rt::Tensor const &logits,
rt::Tensor const &acceptedIndices,
rt::Tensor &gathered,
int32_t batchSize,
int32_t verifyTreeSize,
int32_t maxAcceptDepth,
int32_t vocabSize,
cudaStream_t stream,
)#

Spec-decode verify: gather accepted logit rows from the full verify-tree logits tensor.

For each (batch, depth) pair, copies the logit row at tree position acceptedIndices[batch * maxAcceptDepth + depth] into the output gathered tensor at row batch * maxAcceptDepth + depth. Enables running log-softmax only on the accepted rows (typically 4–7) rather than the full verify tree (60 rows). Used only by speculative-decode verification (EAGLE / MTP) when collecting per-token logprobs.

Parameters:

  • logits – [in] GPU float32 [batchSize * verifyTreeSize, vocabSize]

  • acceptedIndices – [in] GPU int32 [batchSize * maxAcceptDepth] — tree-position indices

  • gathered – [out] GPU float32 [batchSize * maxAcceptDepth, vocabSize]

  • batchSize – [in] Active batch size

  • verifyTreeSize – [in] Rows per batch item in logits

  • maxAcceptDepth – [in] draftingStep + 1

  • vocabSize – [in] Vocabulary size

  • stream – [in] CUDA stream

void trt_edgellm::mapReducedVocabToFullVocab(
rt::Tensor &vocabIds,
rt::Tensor const &vocabMappingTable,
cudaStream_t stream,
)#

Map reduced vocabulary IDs to full vocabulary IDs using a lookup table (in-place).

Performs in-place mapping from reduced vocabulary space to full vocabulary space using the provided mapping table: vocabIds[i] = vocabMappingTable[vocabIds[i]]

The operation is performed in-place, modifying the input tensor directly.

Parameters:

  • vocabIds – [inout] Tensor [GPU, Int32] containing reduced vocabulary IDs as input, will be overwritten with full vocabulary IDs as output

  • vocabMappingTable – [in] Lookup table [GPU, Int32] with shape [reduced_vocab_size] mapping reduced IDs to full IDs

  • stream – [in] CUDA stream to execute the kernel

Throws:

std::runtime_error – If CUDA operations fail

bool trt_edgellm::shouldUseNonGreedySampling(
float temperature,
int64_t topK,
float topP,
) noexcept#

Decide whether sampling parameters require non-greedy decoding.

Greedy decoding is used for the default tuple (temperature=1.0, topK<=1, topP=1.0) as well as near-zero temperatures, which are treated as greedy to avoid unstable softmax.

Parameters:

  • temperature – Sampling temperature

  • topK – Top-k sampling parameter

  • topP – Top-p sampling parameter

Returns:

True when top-k / top-p sampling should be used, false for greedy top-1