Diffusion Gemma Sampling#

struct DiffusionRandomParams#

Stateless RNG parameters for DiffusionGemma sampler kernels.

Public Members

uint64_t seed = {kDefaultDiffusionRandomSeed}#
uint64_t offset = {0}#
struct DiffusionCanvasInitParams#

Scalar parameters for DiffusionGemma canvas initialization.

Public Members

int32_t vocabSize = {0}#
DiffusionRandomParams random = {}#
struct DiffusionCanvasUpdateParams#

Scalar parameters for DiffusionGemma entropy-bound acceptance and re-noising.

Public Members

float entropyThreshold = {0.0F}#
float entropyBound = {0.0F}#
int32_t stabilityWindow = {0}#
bool forceAccept = {false}#
int32_t vocabSize = {0}#
DiffusionRandomParams random = {}#
void trt_edgellm::sampleDiffusionTokensFromLogits(
rt::Tensor const &input,
rt::Tensor &sampledIds,
float temperature,
cudaStream_t stream,
DiffusionRandomParams const &random = DiffusionRandomParams{}
)#

Sample DiffusionGemma canvas token IDs from logits with Gumbel-max.

Samples each row from softmax(logits / temperature) using a stateless Gumbel-max transform. The output tensor may be either flat [rows] or canvas-shaped [batch-size, canvas-len]; its volume must match the input row count.

Parameters:

  • input – [in] Input tensor [GPU, Float/Half/BF16] with shape [rows, vocab-size]

  • sampledIds – [out] Sampled token IDs [GPU, Int32], volume [rows]

  • temperature – [in] Softmax temperature

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

  • random – [in] RNG seed and offset used by the stateless Gumbel sampler

void trt_edgellm::sampleDiffusionTokensAndComputeEntropy(
rt::Tensor const &input,
rt::Tensor &sampledIds,
rt::Tensor &topIndices,
rt::Tensor &entropy,
float temperature,
cudaStream_t stream,
DiffusionRandomParams const &random = DiffusionRandomParams{}
)#

Compute DiffusionGemma Gumbel samples, argmax token IDs, and entropy from logits.

This is equivalent to calling selectArgmaxAndComputeEntropy and sampleDiffusionTokensFromLogits with the same logits and RNG parameters, but it computes argmax, Gumbel sample, and entropy in one vocab pass and removes the separate entropy scan.

Parameters:

  • input – [in] Input tensor [GPU, Float/Half/BF16] with shape [rows, vocab-size]

  • sampledIds – [out] Gumbel-sampled token IDs [GPU, Int32], volume [rows]

  • topIndices – [out] Argmax token IDs [GPU, Int32] with shape [rows, 1]

  • entropy – [out] Entropy values [GPU, Float] with shape [rows]

  • temperature – [in] Softmax temperature

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

  • random – [in] RNG seed and offset used by the stateless Gumbel sampler

void trt_edgellm::selectDiffusionArgmaxFromLogits(
rt::Tensor const &input,
rt::Tensor &topIndices,
float temperature,
cudaStream_t stream
)#

Compute only DiffusionGemma argmax token IDs from logits.

This is used by the final force-accept denoise step where the update policy commits argmax tokens and does not consume sampled IDs or entropy values.

Parameters:

  • input – [in] Input tensor [GPU, Float/Half/BF16] with shape [rows, vocab-size]

  • topIndices – [out] Argmax token IDs [GPU, Int32] with shape [rows, 1]

  • temperature – [in] Positive logit scaling temperature

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

void trt_edgellm::initializeDiffusionCanvas(
rt::Tensor &canvasIds,
rt::Tensor &previousArgmaxIds,
rt::Tensor &stableCounts,
rt::Tensor &acceptedMask,
rt::Tensor &prefixLengths,
DiffusionCanvasInitParams const &params,
cudaStream_t stream
)#

Initialize DiffusionGemma denoise canvas and sampler state on GPU.

Fills the denoise canvas with random token IDs, resets previous argmax IDs, stability counters, accepted mask, and accepted prefix lengths.

Parameters:

  • canvasIds – [out] Randomized canvas token IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • previousArgmaxIds – [out] Previous argmax IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • stableCounts – [out] Consecutive argmax counters including the current observation [GPU, Int32] with shape [batch-size, canvas-len]. Convergence requires this count to be greater than stabilityWindow, matching HF DiffusionGemma history semantics.

  • acceptedMask – [out] Accepted-token mask [GPU, Int8] with shape [batch-size, canvas-len]

  • prefixLengths – [out] Accepted prefix lengths [GPU, Int32] with shape [batch-size]

  • params – [in] Scalar initialization parameters including vocab size and RNG state

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

void trt_edgellm::diffusionSampleAndUpdateCanvas(
rt::Tensor const &sampledIds,
rt::Tensor const &argmaxIds,
rt::Tensor const &entropy,
rt::Tensor &canvasIds,
rt::Tensor &argmaxCanvasIds,
rt::Tensor &previousArgmaxIds,
rt::Tensor &stableCounts,
rt::Tensor &acceptedMask,
rt::Tensor &prefixLengths,
DiffusionCanvasUpdateParams const &params,
cudaStream_t stream,
rt::OptionalInputTensor validCanvasLengths = std::nullopt
)#

Apply DiffusionGemma entropy-bound acceptance and re-noising on GPU.

Consumes Gumbel-sampled token IDs, argmax token IDs, and entropy values. The next denoise canvas keeps Gumbel-sampled IDs for positions selected by the entropy-bound budget and uniformly re-noises rejected positions. The argmax canvas is tracked separately for convergence checks and for the causal commit pass.

Parameters:

  • sampledIds – [in] Gumbel-sampled token IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • argmaxIds – [in] Top-1 token IDs [GPU, Int32] with shape [batch-size * canvas-len, 1]

  • entropy – [in] Entropy values [GPU, Float] with shape [batch-size * canvas-len]

  • canvasIds – [out] Next denoise canvas IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • argmaxCanvasIds – [out] Argmax canvas IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • previousArgmaxIds – [out] Previous argmax IDs [GPU, Int32] with shape [batch-size, canvas-len]

  • stableCounts – [out] Consecutive argmax counters including the current observation [GPU, Int32] with shape [batch-size, canvas-len]. Convergence requires this count to be greater than stabilityWindow, matching HF DiffusionGemma history semantics.

  • acceptedMask – [out] Entropy-bound budget mask [GPU, Int8] with shape [batch-size, canvas-len]

  • prefixLengths – [out] Sticky converged prefix lengths [GPU, Int32] with shape [batch-size]

  • params – [in] Scalar update parameters including entropy policy, force-accept state, vocab size, and RNG state

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

  • validCanvasLengths – [in] Optional per-batch valid lengths [GPU, Int32] with shape [batch-size].

Warning

doxygenfunction: Unable to resolve function “trt_edgellm::compactDiffusionCanvas” with arguments None in doxygen xml output for project “TensorRT Edge-LLM” from directory: ../cpp_docs/xml. Potential matches: 

- void compactDiffusionCanvas(rt::Tensor const &canvasIds, rt::Tensor &commitCanvasIds, int32_t batchSize, int32_t canvasLen, int32_t blockLen, cudaStream_t stream)
- void compactDiffusionCanvas(rt::Tensor const &canvasIds, rt::Tensor const &commitLengths, rt::Tensor &commitCanvasIds, int32_t batchSize, int32_t canvasLen, int32_t maxBlockLen, int32_t padTokenId, cudaStream_t stream)