Embedding Kernels#

void trt_edgellm::kernel::generateMultimodalIndices(
rt::Tensor const &inputIds,
rt::Tensor &multimodalIndices,
std::optional<int32_t> imageTokenId = std::nullopt,
std::optional<int32_t> audioTokenId = std::nullopt,
cudaStream_t stream = nullptr
)#

Generate per-position multimodal indices on-device from GPU token IDs.

For each position holding an image/audio placeholder token, writes the running count of that modality’s placeholders seen so far (the row of imageEmbeds/audioEmbeds to insert); other positions get 0. Counters are global across the whole [batchSize, seqLen] range in batch-major order, matching the host generateMultimodalIndices reference. Runs entirely on stream — no host round-trip, so no D2H/H2D copy is needed to feed embeddingLookup/assembleDeepstackEmbedding.

Parameters:

  • inputIds – [in] GPU token IDs [batchSize, seqLen] (INT32)

  • multimodalIndices – [out] GPU indices [batchSize, seqLen] (INT32), same element count as inputIds

  • imageTokenId – [in] Image placeholder token id, or std::nullopt if no image

  • audioTokenId – [in] Audio placeholder token id, or std::nullopt if no audio

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::assembleDeepstackEmbedding(
rt::Tensor const &inputIds,
rt::Tensor const &deepstackFeatures,
rt::Tensor &deepstackEmbeds,
cudaStream_t stream,
int32_t imageTokenId = 0,
rt::OptionalInputTensor multimodalIndices = std::nullopt
)#

Assemble deepstack embeddings by extracting image token embeddings from deepstack features.

This function processes input token IDs and selectively extracts embeddings for image tokens from the provided deepstack features. Image tokens are identified by the explicit imageTokenId, and multimodalIndices selects the deepstack feature row for each image position.

Parameters:

  • inputIds – [in] Input token IDs with shape [batchSize, seqLen]

  • deepstackFeatures – [in] Deepstack image features with shape [numImageTokens, hiddenSize]

  • imageTokenId – [in] Image token ID; positions with this id receive a deepstack feature row

  • multimodalIndices – [in] Pre-computed indices for image embeddings [batchSize, seqLen]

  • deepstackEmbeds – [out] Output embeddings with shape [batchSize, seqLen, hiddenSize]

  • stream – [in] CUDA stream for execution

Throws:

std::runtime_error – if tensor shapes or data types are invalid

void trt_edgellm::kernel::embeddingLookup(
rt::Tensor const &inputIds,
rt::Tensor const &embeddingTable,
rt::OptionalInputTensor scales,
rt::Tensor &output,
cudaStream_t stream,
rt::OptionalInputTensor multimodalIndices = std::nullopt,
std::optional<int32_t> imageTokenId = std::nullopt,
rt::OptionalInputTensor imageEmbeds = std::nullopt,
std::optional<int32_t> audioTokenId = std::nullopt,
rt::OptionalInputTensor audioEmbeds = std::nullopt
)#

Embedding lookup for all modalities (supports FP16 and FP8 tables).

Produces input embeddings for a batch of token ids. Text tokens are looked up from the embedding table. When image and/or audio embeddings are supplied (prefill), positions whose token id equals imageTokenId / audioTokenId are filled from imageEmbeds / audioEmbeds, selected by multimodalIndices; with no image/audio inputs (decode) it performs a pure text lookup. Automatically dispatches to the FP16 or FP8 implementation based on the embedding table’s datatype; FP8 tables require scales for per-group dequantization.

Note

imageTokenId and audioTokenId must differ when both are provided.

Parameters:

  • inputIds – [in] Input token IDs with shape [batchSize, seqLen]

  • embeddingTable – [in] Text embedding table with shape [vocabSize, hiddenSize] (FP16 or FP8)

  • scales – [in] FP32 per-group scales with shape [vocabSize, hiddenSize / blockSize]; required when embeddingTable is FP8, std::nullopt for FP16

  • output – [out] Hidden states with shape [batchSize, seqLen, hiddenSize]

  • stream – [in] CUDA stream for execution

  • multimodalIndices – [in] Per-position indices into imageEmbeds/audioEmbeds [batchSize, seqLen]; required when image or audio inputs are provided, std::nullopt otherwise

  • imageTokenId – [in] Token ID marking image positions, or std::nullopt if no image input

  • imageEmbeds – [in] Image embeddings [totalImageTokens, hiddenSize] (FP16), or std::nullopt

  • audioTokenId – [in] Token ID marking audio positions, or std::nullopt if no audio input

  • audioEmbeds – [in] Audio embeddings [totalAudioTokens, hiddenSize] (FP16), or std::nullopt

Throws:

std::runtime_error – if tensor shapes or data types are invalid, or FP8 is not supported

void trt_edgellm::kernel::gemma4PleGather(
rt::Tensor const &inputIds,
rt::Tensor const &pleTable,
rt::Tensor &outputBuffer,
int32_t numLayers,
int32_t pleHiddenSize,
int32_t imageTokenId,
int32_t audioTokenId,
cudaStream_t stream
)#

Gather Gemma4 per-layer token-identity embeddings.

Parameters:

  • inputIds – [in] Input token IDs with shape [batchSize, seqLen]

  • pleTable – [in] PLE table with shape [vocabSize, numLayers * pleHiddenSize]

  • outputBuffer – [inout] Backing tensor for all per-layer outputs; shape [numLayers, maxBatch, maxSeq, hidden]

  • numLayers – [in] Number of PLE layer outputs

  • pleHiddenSize – [in] Hidden size of each PLE output

  • imageTokenId – [in] Optional image token ID to zero-fill (-1 = unused)

  • audioTokenId – [in] Optional audio token ID to zero-fill (-1 = unused)

  • stream – [in] CUDA stream for execution