Talker Mlp Kernels#

struct PrefillRowDesc#

One prefill output row = srcA (+ srcB when non-null), both device pointers to [H] rows.

The per-row math is identical to invokeAssistantPreamble (vectorized half2 add, same order), so rows described identically produce bit-identical output.

Public Members

half const *srcA#

Required source row [hiddenDim].

half const *srcB#

Optional addend row [hiddenDim]; nullptr = copy srcA only.

void trt_edgellm::kernel::invokeTalkerMLP(
rt::Tensor const &input,
rt::Tensor const &fc1Weight,
rt::Tensor const &fc1Bias,
rt::Tensor const &fc2Weight,
rt::Tensor const &fc2Bias,
rt::Tensor &output,
rt::Tensor &workspace,
cudaStream_t stream
)#

Two-layer MLP with SiLU activation (Talker projection layers)

Performs: output = FC2(SiLU(FC1(input) + bias1)) + bias2 Where FC1: [inputDim → hiddenDim], FC2: [hiddenDim → outputDim]

Architecture: input [N, 2048] ↓ FC1 (Linear) [N, 2048] + bias1 ↓ SiLU [N, 2048] ↓ FC2 (Linear) [N, 1024] + bias2 ↓ output [N, 1024]

GEMM is performed via CuTe DSL compiled kernels (FP16 I/O, FP32 accumulation).

Parameters:

  • input – [in] Input tensor with shape [numTokens, inputDim] (FP16)

  • fc1Weight – [in] FC1 weight matrix with shape [hiddenDim, inputDim] (FP16, row-major)

  • fc1Bias – [in] FC1 bias vector with shape [hiddenDim] (FP16)

  • fc2Weight – [in] FC2 weight matrix with shape [outputDim, hiddenDim] (FP16, row-major)

  • fc2Bias – [in] FC2 bias vector with shape [outputDim] (FP16)

  • output – [out] Output tensor with shape [numTokens, outputDim] (FP16)

  • workspace – [inout] Workspace buffer for intermediate FC1 output [numTokens, hiddenDim] (FP16)

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::invokeLinearLayer(
rt::Tensor const &input,
rt::Tensor const &weight,
rt::Tensor const &bias,
rt::Tensor &output,
cudaStream_t stream
)#

Single linear layer: output = input @ weight.T + bias.

GEMM is performed via CuTe DSL compiled kernels (FP16 I/O, FP32 accumulation).

Parameters:

  • input – [in] Input tensor with shape [N, inputDim] (FP16)

  • weight – [in] Weight matrix with shape [outputDim, inputDim] (FP16, row-major)

  • bias – [in] Bias vector with shape [outputDim] (FP16)

  • output – [out] Output tensor with shape [N, outputDim] (FP16)

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::invokeGather(
rt::Tensor const &source,
rt::Tensor const &indices,
rt::Tensor &output,
cudaStream_t stream
)#

Gather operation: select rows from source tensor by indices.

Performs: output[i] = source[indices[i]] where each row has hiddenDim elements.

Parameters:

  • source – [in] Source tensor with shape [srcNumTokens, hiddenDim] (FP16)

  • indices – [in] Indices tensor with shape [numIndices] (INT32)

  • output – [out] Output tensor with shape [numIndices, hiddenDim] (FP16)

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::invokeScatter(
rt::Tensor const &source,
rt::Tensor const &indices,
rt::Tensor &output,
cudaStream_t stream
)#

Scatter operation: place rows from source to output by indices.

Performs: output[indices[i]] = source[i] where each row has hiddenDim elements.

Parameters:

  • source – [in] Source tensor with shape [numIndices, hiddenDim] (FP16)

  • indices – [in] Indices tensor with shape [numIndices] (INT32)

  • output – [out] Output tensor with shape [dstNumTokens, hiddenDim] (FP16)

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::invokeAssistantPreamble(
rt::Tensor const &projected,
rt::Tensor const &ttsPadEmbed,
rt::Tensor const &ttsBosEmbed,
rt::Tensor const &ttsEosEmbed,
rt::Tensor const &talkerEmbTable,
int32_t codecNothinkId,
int32_t codecThinkBosId,
int32_t codecThinkEosId,
int32_t speakerId,
int32_t codecPadId,
int32_t codecBosId,
int32_t codecThinkId,
int32_t languageId,
int32_t textLen,
rt::Tensor &output,
cudaStream_t stream
)#

Fused non-streaming assistant preamble construction for TTS input projection.

Builds the complete non-streaming prefill buffer in one pass. Total rows written = P + textLen + 2, where P = 8 without language conditioning or P = 9 when languageId >= 0 (CustomVoice language conditioning inserts one extra row).

Row layout without language (P = 8, byte-identical to the historical layout):

[3]: ttsPadEmbed + talkerEmbTable[codecNothinkId] [4]: ttsPadEmbed + talkerEmbTable[codecThinkBosId] [5]: ttsPadEmbed + talkerEmbTable[codecThinkEosId] [6]: ttsPadEmbed + talkerEmbTable[speakerId] [7]: ttsBosEmbed + talkerEmbTable[codecPadId]

Row layout with language (P = 9; think-token replaces the no-think token, language row inserted):

[3]: ttsPadEmbed + talkerEmbTable[codecThinkId] [4]: ttsPadEmbed + talkerEmbTable[codecThinkBosId] [5]: ttsPadEmbed + talkerEmbTable[languageId] [6]: ttsPadEmbed + talkerEmbTable[codecThinkEosId] [7]: ttsPadEmbed + talkerEmbTable[speakerId] [8]: ttsBosEmbed + talkerEmbTable[codecPadId]

Shared text/suffix rows: [P..P+N-2]: projected[3+i] + talkerEmbTable[codecPadId] (text tokens, N=textLen) [P+N-1]: projected[3+N-1] + talkerEmbTable[codecBosId] (last text = start-of-generation) [P+N]: ttsEosEmbed + talkerEmbTable[codecPadId] [P+N+1]: ttsPadEmbed + talkerEmbTable[codecBosId]

Parameters:

  • projected – MLP output [seqLen, H] (FP16)

  • ttsPadEmbed/ttsBosEmbed/ttsEosEmbed – TTS special embeddings [H] (FP16)

  • talkerEmbTable – Talker embedding table [vocabSize, H] (FP16)

  • codecNothinkId..codecBosId – Codec token IDs used in the prefix/suffix rows

  • speakerId – Speaker codec token ID

  • codecThinkId – Codec think token ID (used instead of codecNothinkId when languageId >= 0; ignored otherwise, pass -1 if unavailable)

  • languageId – Language codec token ID; -1 disables language conditioning (8-row prefix)

  • textLen – Number of text token rows (N)

  • output – Full output buffer [P+N+2, H] (FP16)

  • stream – CUDA stream

void trt_edgellm::kernel::invokePrefillRowAssemble(
PrefillRowDesc const *deviceDescs,
int32_t numRows,
int32_t hiddenDim,
rt::Tensor &output,
cudaStream_t stream
)#

Descriptor-driven prefill row assembly (generalization of invokeAssistantPreamble)

Assembles arbitrary prefill layouts (instruction segments, no-speaker VoiceDesign prefixes, continuous speaker embeddings, ICL segments) from a host-built row descriptor list. One block per row; each row is srcA (+ srcB) with the same vectorized half2 add as invokeAssistantPreamble.

Parameters:

  • deviceDescs – Device array of numRows descriptors (uploaded by the caller)

  • numRows – Number of output rows

  • hiddenDim – Row width (must be a multiple of 8)

  • output – Output buffer [numRows, hiddenDim] (FP16)

  • stream – CUDA stream

void trt_edgellm::kernel::invokeSumCodecEmbeddings(
int64_t const *refCodes,
half const *const *tablePtrs,
int32_t numFrames,
int32_t numGroups,
int32_t hiddenDim,
rt::Tensor &output,
cudaStream_t stream
)#

Sum per-frame codec embeddings across all code groups (voice-clone ICL prompt)

For each reference frame t: output[t] = sum_{g=0}^{numGroups-1} tables[g][codes[t][g]]. Group 0 uses the Talker codec embedding table; groups 1..numGroups-1 the CodePredictor tables — the caller passes one device pointer per group. Accumulation in FP32.

Parameters:

  • refCodes – Device codes [numFrames, numGroups] (INT32)

  • tablePtrs – Device array of numGroups table pointers, each [vocab, hiddenDim] (FP16)

  • numFrames – Reference frame count

  • numGroups – Code group count (16 for the 12Hz tokenizer)

  • hiddenDim – Embedding width

  • output – Output [numFrames, hiddenDim] (FP16)

  • stream – CUDA stream refCodes is INT64 so the codec-encoder engine output is consumed in place (no host round-trip or dtype conversion).

void trt_edgellm::kernel::invokeCastFp32ToFp16(
float const *input,
half *output,
int64_t numElements,
cudaStream_t stream
)#

Elementwise FP32 -> FP16 cast on device (small utility for engine-output adaptation)

void trt_edgellm::kernel::invokeResidualConnection(
rt::Tensor const &codecHiddens,
rt::Tensor const &embTable0,
rt::Tensor const &embTableLast,
int32_t code0,
int32_t codeLast,
half const *addend,
rt::Tensor &output,
cudaStream_t stream
)#

Fused residual connection for TTS decode input.

Computes: output = embed0[code0] + embedLast[codeLast] + addend + sum(codecHiddens[1..N-1]) where N = numCodesPerFrame (inferred from codecHiddens shape). Eliminates 7 separate dispatches (2x H→D, 2x embLookup, 2x D→D, sumReduce) in one kernel.

Parameters:

  • codecHiddens – [1, numCodesPerFrame, H] buffer — inner rows pre-filled by CodePredictor (FP16)

  • embTable0 – Talker embedding table [vocabSize, H] (FP16) — for embed(code0)

  • embTableLast – CodePredictor embedding table[-1] [vocabSize, H] (FP16) — for embed(codeLast)

  • code0/codeLast – Token IDs passed as scalars (no H→D upload needed)

  • addend – Row pointer [H] — trailing_text_hidden[generationStep] or tts_pad_embed (FP16)

  • output – Output tensor [1, 1, H] (FP16)

  • stream – CUDA stream

void trt_edgellm::kernel::invokeTalkerLogitAdjust(
rt::Tensor const &seenTokens,
rt::Tensor &logits,
int32_t suppressStart,
int32_t suppressEnd,
int32_t codecEosId,
int32_t numSeenTokens,
float repetitionPenalty,
cudaStream_t stream
)#

Adjust Talker logits: suppress special tokens and apply repetition penalty.

Performs two in-place modifications on the logits before sampling:

  1. Suppression: sets logits[i] = -inf for all i in [suppressStart, suppressEnd), except for codecEosId which is always preserved.

  2. Repetition penalty: for each token in seenTokens[], divides positive logits by repetitionPenalty and multiplies negative logits by repetitionPenalty, matching the HuggingFace repetition_penalty convention.

Operates on FP32 logits tensor with shape [1, vocabSize].

Parameters:

  • seenTokens – [in] GPU tensor of previously generated token IDs [maxAudioLength] INT32

  • logits – [inout] Logits tensor [1, vocabSize] (FP32, in-place)

  • suppressStart – [in] Start of suppress range (inclusive)

  • suppressEnd – [in] End of suppress range (exclusive)

  • codecEosId – [in] Token ID exempt from suppression (EOS must remain samplable)

  • numSeenTokens – [in] Number of valid entries in seenTokens (0 to disable penalty)

  • repetitionPenalty – [in] Penalty factor >= 1.0 (1.0 = no penalty)

  • stream – [in] CUDA stream for execution

void trt_edgellm::kernel::invokeSpeakerCodecSum(
rt::Tensor const &codes,
rt::Tensor const &embPtrTable,
rt::Tensor const &embVocabSizes,
rt::Tensor &output,
cudaStream_t stream
)#

Per-position sum across codec groups: out[d] = sum_g embPtrTable[g][codes[g]][d]; codes[g] < 0 skips group g. embPtrTable is INT8-typed to carry __half const*[numCodeGroups] (rt::Tensor has no pointer-array dtype). numCodeGroups = codes.shape[0]; hiddenSize = output.volume().